Problem p03453 (python)

Fixed submission:
def main():
    import sys
    sys.setrecursionlimit(100000)
    input = sys.stdin.readline
    from heapq import heappop, heappush
    mod = 10**9+7
    INF = 1<<60
    N, M = map(int, input().split())
    S, T = map(int, input().split())
    E = [[] for _ in range(N+1)]
    edges = []
    # for _ in range(M):
    #     u, v, d = map(int, input().split())
    #     E[u].append((v, d))
    #     E[v].append((u, d))
    #     edges.append((u, v, d))
    for u, v, d in zip(*[iter(map(int, sys.stdin.read().split()))]*3):
        E[u].append((v, d))
        E[v].append((u, d))
        edges.append((u, v, d))

    def dijksrtra(start):
        Dist = [INF] * (N+1)
        Dist[start] = 0
        q = [0<<17 | start]
        mask = (1<<17)-1
        while q:
            dist_v = heappop(q)
            v = dist_v & mask
            dist = dist_v >> 17
            if Dist[v] != dist:
                continue
            for u, d in E[v]:
                new_dist = dist + d
                if Dist[u] > new_dist:
                    Dist[u] = new_dist
                    heappush(q, new_dist<<17 | u)
        return Dist

    Dist_S = dijksrtra(S)
    Dist_T = dijksrtra(T)
    dist_st = Dist_S[T]

    DAG_edges = []  # S -> T
    DAG = [[] for _ in range(N+1)]
    DAG_rev = [[] for _ in range(N+1)]
    for u, v, d in edges:
        if Dist_S[u] + Dist_T[v] + d == dist_st:
            DAG_edges.append((u, v))
            DAG[u].append(v)
            DAG_rev[v].append(u)
        elif Dist_T[u] + Dist_S[v] + d == dist_st:
            DAG_edges.append((v, u))
            DAG[v].append(u)
            DAG_rev[u].append(v)


    # トポロジカルソート
    V = []
    rem = [0] * (N+1)
    for _, v in DAG_edges:
        rem[v] += 1
    q = [S]
    while q:
        v = q.pop()
        V.append(v)
        for u in DAG[v]:
            rem[u] -= 1
            if rem[u] == 0:
                q.append(u)


    # n_paths_S = [-1] * (N+1)
    # n_paths_T = [-1] * (N+1)
    n_paths_S = [0] * (N+1)
    n_paths_T = [0] * (N+1)
    n_paths_S[S] = 1
    n_paths_T[T] = 1
    # def calc_n_paths_S(v):
    #     if n_paths_S[v] != -1:
    #         return n_paths_S[v]
    #     res = 0
    #     for u in DAG_rev[v]:
    #         res = (res + calc_n_paths_S(u)) % mod
    #     n_paths_S[v] = res
    #     return res
    # def calc_n_paths_T(v):
    #     if n_paths_T[v] != -1:
    #         return n_paths_T[v]
    #     res = 0
    #     for u in DAG[v]:
    #         res = (res + calc_n_paths_T(u)) % mod
    #     n_paths_T[v] = res
    #     return res

    # ans = calc_n_paths_S(T)  # 全経路数
    # calc_n_paths_T(S)

    for v in V:
        n = n_paths_S[v]
        for u in DAG[v]:
            n_paths_S[u] += n
    for v in V[::-1]:
        n = n_paths_T[v]
        for u in DAG_rev[v]:
            n_paths_T[u] += n
    ans = n_paths_S[T]
    ans = ans * ans % mod

    for v, u in DAG_edges:  # 辺ですれ違う場合
        if Dist_S[v]*2 < dist_st and dist_st < Dist_S[u]*2:
            ans = (ans - (n_paths_S[v]*n_paths_T[u])**2) % mod

    # 頂点ですれ違う場合
    for v, (dist, ns, nt) in enumerate(zip(Dist_S, n_paths_S, n_paths_T)):
        if dist*2 == dist_st:
            ans = (ans - (ns*nt)**2) % mod

    print(ans)

main()

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
from heapq import heappop, heappush
mod = 10**9+7
INF = 1<<60
N, M = map(int, input().split())
S, T = map(int, input().split())
E = [[] for _ in range(N+1)]
edges = []
for _ in range(M):
    u, v, d = map(int, input().split())
    E[u].append((v, d))
    E[v].append((u, d))
    edges.append((u, v, d))

def dijksrtra(start):
    Dist = [INF] * (N+1)
    Dist[start] = 0
    q = [(0, start)]
    while q:
        dist, v = heappop(q)
        if Dist[v] != dist:
            continue
        for u, d in E[v]:
            new_dist = dist + d
            if Dist[u] > new_dist:
                Dist[u] = new_dist
                heappush(q, (new_dist, u))
    return Dist

Dist_S = dijksrtra(S)
Dist_T = dijksrtra(T)
dist_st = Dist_S[T]

DAG_edges = []  # S -> T
DAG = [[] for _ in range(N+1)]
DAG_rev = [[] for _ in range(N+1)]
for u, v, d in edges:
    if Dist_S[u] + Dist_T[v] + d == dist_st:
        DAG_edges.append((u, v, d))
        DAG[u].append((v, d))
        DAG_rev[v].append((u, d))
    elif Dist_T[u] + Dist_S[v] + d == dist_st:
        DAG_edges.append((v, u, d))
        DAG[v].append((u, d))
        DAG_rev[u].append((v, d))

n_paths_S = [-1] * (N+1)
n_paths_T = [-1] * (N+1)
n_paths_S[S] = 1
n_paths_T[T] = 1
def calc_n_paths_S(v):
    if n_paths_S[v] != -1:
        return n_paths_S[v]
    res = 0
    for u, _ in DAG_rev[v]:
        res = (res + calc_n_paths_S(u)) % mod
    n_paths_S[v] = res
    return res
def calc_n_paths_T(v):
    if n_paths_T[v] != -1:
        return n_paths_T[v]
    res = 0
    for u, _ in DAG[v]:
        res = (res + calc_n_paths_T(u)) % mod
    n_paths_T[v] = res
    return res

ans = calc_n_paths_S(T)  # 全経路数
calc_n_paths_T(S)
#assert ans == calc_n_paths_T(S), (n_paths_S, n_paths_T)

ans = ans * ans % mod

for v, u, _ in DAG_edges:  # 辺ですれ違う場合
    if Dist_S[v]*2 < dist_st and dist_st < Dist_S[u]*2:
        ans = (ans - (n_paths_S[v]*n_paths_T[u])**2) % mod

# 頂点ですれ違う場合
for v, (dist, ns, nt) in enumerate(zip(Dist_S, n_paths_S, n_paths_T)):
    if dist*2 == dist_st:
        ans = (ans - (ns*nt)**2) % mod

print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
from heapq import heappop, heappush
mod = 10**9+7
INF = 1<<60
N, M = map(int, input().split())
S, T = map(int, input().split())
E = [[] for _ in range(N+1)]
edges = []
for _ in range(M):
    u, v, d = map(int, input().split())
    E[u].append((v, d))
    E[v].append((u, d))
    edges.append((u, v, d))

def dijksrtra(start):
    Dist = [INF] * (N+1)
    Dist[start] = 0
    q = [(0, start)]
    while q:
        dist, v = heappop(q)
        if Dist[v] != dist:
            continue
        for u, d in E[v]:
            new_dist = dist + d
            if Dist[u] > new_dist:
                Dist[u] = new_dist
                heappush(q, (new_dist, u))
    return Dist

Dist_S = dijksrtra(S)
Dist_T = dijksrtra(T)
dist_st = Dist_S[T]

DAG_edges = []  # S -> T
DAG = [[] for _ in range(N+1)]
DAG_rev = [[] for _ in range(N+1)]
for u, v, d in edges:
    if Dist_S[u] + Dist_T[v] + d == dist_st:
        DAG_edges.append((u, v, d))
        DAG[u].append((v, d))
        DAG_rev[v].append((u, d))
    elif Dist_T[u] + Dist_S[v] + d == dist_st:
        DAG_edges.append((v, u, d))
        DAG[v].append((u, d))
        DAG_rev[u].append((v, d))

n_paths_S = [-1] * (N+1)
n_paths_T = [-1] * (N+1)
n_paths_S[S] = 1
n_paths_T[T] = 1
def calc_n_paths_S(v):
    if n_paths_S[v] != -1:
        return n_paths_S[v]
    res = 0
    for u, _ in DAG_rev[v]:
        res = (res + calc_n_paths_S(u)) % mod
    n_paths_S[v] = res
    return res
def calc_n_paths_T(v):
    if n_paths_T[v] != -1:
        return n_paths_T[v]
    res = 0
    for u, _ in DAG[v]:
        res = (res + calc_n_paths_T(u)) % mod
    n_paths_T[v] = res
    return res

ans = calc_n_paths_S(T)  # 全経路数
assert ans == calc_n_paths_T(S), (n_paths_S, n_paths_T)

ans = ans * ans % mod

for v, u, _ in DAG_edges:  # 辺ですれ違う場合
    if Dist_S[v]*2 < dist_st and dist_st < Dist_S[u]*2:
        ans = (ans - (n_paths_S[v]*n_paths_T[u])**2) % mod

# 頂点ですれ違う場合
for v, (dist, ns, nt) in enumerate(zip(Dist_S, n_paths_S, n_paths_T)):
    if dist*2 == dist_st:
        ans = (ans - (ns*nt)**2) % mod

print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
from heapq import heappop, heappush
mod = 10**9+7
INF = 1<<60
N, M = map(int, input().split())
S, T = map(int, input().split())
E = [[] for _ in range(N+1)]
edges = []
for _ in range(M):
    u, v, d = map(int, input().split())
    E[u].append((v, d))
    E[v].append((u, d))
    edges.append((u, v, d))

def dijksrtra(start):
    Dist = [INF] * (N+1)
    Dist[start] = 0
    q = [(0, start)]
    while q:
        dist, v = heappop(q)
        if Dist[v] != dist:
            continue
        for u, d in E[v]:
            new_dist = dist + d
            if Dist[u] > new_dist:
                Dist[u] = new_dist
                heappush(q, (new_dist, u))
    return Dist

Dist_S = dijksrtra(S)
Dist_T = dijksrtra(T)
dist_st = Dist_S[T]

DAG_edges = []  # S -> T
DAG = [[] for _ in range(N+1)]
DAG_rev = [[] for _ in range(N+1)]
for u, v, d in edges:
    if Dist_S[u] + Dist_T[v] + d == dist_st:
        DAG_edges.append((u, v, d))
        DAG[u].append((v, d))
        DAG_rev[v].append((u, d))
    elif Dist_T[u] + Dist_S[v] + d == dist_st:
        DAG_edges.append((v, u, d))
        DAG[v].append((u, d))
        DAG_rev[u].append((v, d))

n_paths_S = [-1] * (N+1)
n_paths_T = [-1] * (N+1)
n_paths_S[S] = 1
n_paths_T[T] = 1
def calc_n_paths_S(v):
    if n_paths_S[v] != -1:
        return n_paths_S[v]
    res = 0
    for u, _ in DAG_rev[v]:
        res = (res + calc_n_paths_S(u)) % mod
    n_paths_S[v] = res
    return res
def calc_n_paths_T(v):
    if n_paths_T[v] != -1:
        return n_paths_T[v]
    res = 0
    for u, _ in DAG[v]:
        res = (res + calc_n_paths_T(u)) % mod
    n_paths_T[v] = res
    return res

ans = calc_n_paths_S(T)  # 全経路数
assert ans == calc_n_paths_T(S), (n_paths_S, n_paths_T)

ans = ans * ans % mod

half = dist_st / 2.0
assert half * 2 == dist_st
for v, u, _ in DAG_edges:  # 辺ですれ違う場合
    if Dist_S[v] < half and half < Dist_S[u]:
        ans = (ans - (n_paths_S[v]*n_paths_T[u])**2) % mod

# 頂点ですれ違う場合
for v, (dist, ns, nt) in enumerate(zip(Dist_S, n_paths_S, n_paths_T)):
    if dist == half:
        ans = (ans - (ns*nt)**2) % mod

print(ans)

Buggy submission (verdict: Runtime Error):
from heapq import heappop, heappush
mod = 10**9+7
N, M = map(int, input().split())
S, T = map(int, input().split())
E = [[] for _ in range(N+1)]
edges = []
for _ in range(M):
    u, v, d = map(int, input().split())
    E[u].append((v, d))
    E[v].append((u, d))
    edges.append((u, v, d))

def dijksrtra(start):
    Dist = [float("inf")] * (N+1)
    Dist[start] = 0
    q = [(0, start)]
    while q:
        dist, v = heappop(q)
        if Dist[v] != dist:
            continue
        for u, d in E[v]:
            new_dist = dist + d
            if Dist[u] > new_dist:
                Dist[u] = new_dist
                heappush(q, (new_dist, u))
    return Dist

Dist_S = dijksrtra(S)
Dist_T = dijksrtra(T)
dist_st = Dist_S[T]

DAG_edges = []  # S -> T
DAG = [[] for _ in range(N+1)]
DAG_rev = [[] for _ in range(N+1)]
for u, v, d in edges:
    if Dist_S[u] + Dist_T[v] + d == dist_st:
        DAG_edges.append((u, v, d))
        DAG[u].append((v, d))
        DAG_rev[v].append((u, d))
    elif Dist_T[u] + Dist_S[v] + d == dist_st:
        DAG_edges.append((v, u, d))
        DAG[v].append((u, d))
        DAG_rev[u].append((v, d))

n_paths_S = [-1] * (N+1)
n_paths_T = [-1] * (N+1)
n_paths_S[S] = 1
n_paths_T[T] = 1
def calc_n_paths_S(v):
    if n_paths_S[v] != -1:
        return n_paths_S[v]
    res = 0
    for u, _ in DAG_rev[v]:
        res = (res + calc_n_paths_S(u)) % mod
    n_paths_S[v] = res
    return res
def calc_n_paths_T(v):
    if n_paths_T[v] != -1:
        return n_paths_T[v]
    res = 0
    for u, _ in DAG[v]:
        res = (res + calc_n_paths_T(u)) % mod
    n_paths_T[v] = res
    return res

ans = calc_n_paths_S(T)  # 全経路数
assert ans == calc_n_paths_T(S), (n_paths_S, n_paths_T)

ans = ans * ans % mod

half = dist_st / 2.0
assert half * 2 == dist_st
for v, u, _ in DAG_edges:  # 辺ですれ違う場合
    if Dist_S[v] < half and half < Dist_S[u]:
        ans = (ans - (n_paths_S[v]*n_paths_T[u])**2) % mod

# 頂点ですれ違う場合
for v, (dist, ns, nt) in enumerate(zip(Dist_S, n_paths_S, n_paths_T)):
    if dist == half:
        ans = (ans - (ns*nt)**2) % mod

print(ans)

Buggy submission (verdict: Runtime Error):
from heapq import heappop, heappush
mod = 10**9+7
N, M = map(int, input().split())
S, T = map(int, input().split())
E = [[] for _ in range(N+1)]
edges = []
for _ in range(M):
    u, v, d = map(int, input().split())
    E[u].append((v, d))
    E[v].append((u, d))
    edges.append((u, v, d))

def dijksrtra(start):
    Dist = [float("inf")] * (N+1)
    Dist[start] = 0
    q = [(0, start)]
    while q:
        dist, v = heappop(q)
        if Dist[v] != dist:
            continue
        for u, d in E[v]:
            new_dist = dist + d
            if Dist[u] > new_dist:
                Dist[u] = new_dist
                heappush(q, (new_dist, u))
    return Dist

Dist_S = dijksrtra(S)
Dist_T = dijksrtra(T)
dist_st = Dist_S[T]

DAG_edges = []  # S -> T
DAG = [[] for _ in range(N+1)]
DAG_rev = [[] for _ in range(N+1)]
for u, v, d in edges:
    if Dist_S[u] + Dist_T[v] + d == dist_st:
        DAG_edges.append((u, v, d))
        DAG[u].append((v, d))
        DAG_rev[v].append((u, d))
    elif Dist_T[u] + Dist_S[v] + d == dist_st:
        DAG_edges.append((v, u, d))
        DAG[v].append((u, d))
        DAG_rev[u].append((v, d))

n_paths_S = [-1] * (N+1)
n_paths_T = [-1] * (N+1)
n_paths_S[S] = 1
n_paths_T[T] = 1
def calc_n_paths_S(v):
    if n_paths_S[v] != -1:
        return n_paths_S[v]
    res = 0
    for u, _ in DAG_rev[v]:
        res = (res + calc_n_paths_S(u)) % mod
    n_paths_S[v] = res
    return res
def calc_n_paths_T(v):
    if n_paths_T[v] != -1:
        return n_paths_T[v]
    res = 0
    for u, _ in DAG[v]:
        res = (res + calc_n_paths_T(u)) % mod
    n_paths_T[v] = res
    return res

ans = calc_n_paths_S(T)  # 全経路数
assert ans == calc_n_paths_T(S), (n_paths_S, n_paths_T)

ans = ans * ans % mod

half = dist_st / 2.0
assert half * 2 == dist_st
for v, u, _ in DAG_edges:  # 辺ですれ違う場合
    if Dist_S[v] < half and half < Dist_S[u]:
        ans = (ans - (n_paths_S[v]*n_paths_T[u])**2) % mod

for v, (dist, ns, nt) in enumerate(zip(Dist_S, n_paths_S, n_paths_T)):
    # 頂点ですれ違う場合
    if dist == half:
        ans = (ans - (ns*nt)**2) % mod

print(ans)