Problem p02986 (python)

Fixed submission:
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)

N, Q = map(int, input().split())

path = [[] for _ in range(N)]

for _ in range(N-1) :
    a, b, c, d = (int(i) for i in input().split())
    path[a-1].append((b-1, c-1, d))
    path[b-1].append((a-1, c-1, d))

# doublingに必要なKを求める
for K in range(18) :
    if 2 ** K >= N :
        break

# dfs
parent = [[-1] * N for _ in range(K)]
rank = [-1 for _ in range(N)]

rank[0] = 0
queue = [0]

while queue :
    cur = queue.pop()
    for nex, _, _ in path[cur] :
        if rank[nex] < 0 :
            queue.append(nex)
            parent[0][nex] = cur
            rank[nex] = rank[cur] + 1

# doubling        
for i in range(1, K) :
    for j in range(N) :
        if parent[i-1][j] > 0 :
            parent[i][j] = parent[i-1][parent[i-1][j]]

# lca
def lca(a, b) :
    if rank[a] > rank[b] :
        a, b = b, a

    diff = rank[b] - rank[a]
    i = 0
    while diff > 0 :
        if diff & 1 :
            b = parent[i][b]
        diff >>= 1
        i += 1
        
    if a == b :
        return a

    for i in range(K-1, -1, -1) :
        if parent[i][a] != parent[i][b] :
            a = parent[i][a]
            b = parent[i][b]

    return parent[0][a]

# Queryの先読み
schedule = [[] for _ in range(N)]
for i in range(Q) : 
    x, y, u, v = map(int, input().split())
    x, u, v = x-1, u-1, v-1
    l = lca(u, v)
    schedule[u].append((i, 1, x, y))
    schedule[v].append((i, 1, x, y))
    schedule[l].append((i, -2, x, y))

ret = [0] * Q
C = [0] * (N-1)
D = [0] * (N-1)

def dfs(cur, pre, tot) :
    for i, t, c, d in schedule[cur] :
        ret[i] += t * (tot - D[c] + C[c] * d)
    
    for nex, c, d in path[cur] :
        if nex == pre :
            continue
        C[c] += 1
        D[c] += d
        dfs(nex, cur, tot + d)
        C[c] -= 1
        D[c] -= d

dfs(0, -1, 0)
   
for i in range(Q) :
    print(ret[i])

Buggy submission (verdict: Runtime Error):
from collections import deque

N, Q = map(int, input().split())

path = [[] for _ in range(N)]

for _ in range(N-1) :
    a, b, c, d = (int(i) for i in input().split())
    path[a-1].append((b-1, c-1, d))
    path[b-1].append((a-1, c-1, d))

# doublingに必要なKを求める
for K in range(18) :
    if 2 ** K >= N :
        break

# bfs
parent = [[-1] * N for i in range(K)]
rank = [-1 for _ in range(N)]
dist = [0 for _ in range(N)]

rank[0] = 0
queue = deque([0])

while queue :
    cur = queue.popleft()
    for nex, _, d in path[cur] :
        if rank[nex] < 0 :
            queue.append(nex)
            parent[0][nex] = cur
            rank[nex] = rank[cur] + 1
            dist[nex] = dist[cur] + d

# doubling        
for i in range(1, K) :
    for j in range(N) :
        if parent[i-1][j] > 0 :
            parent[i][j] = parent[i-1][parent[i-1][j]]

# lca
def lca(a, b) :
    if rank[a] > rank[b] :
        a, b = b, a

    diff = rank[b] - rank[a]
    i = 0
    while diff > 0 :
        if diff & 1 :
            b = parent[i][b]
        diff >>= 1
        i += 1
        
    if a == b :
        return a

    for i in range(K-1, -1, -1) :
        if parent[i][a] != parent[i][b] :
            a = parent[i][a]
            b = parent[i][b]

    return parent[0][a]

# Queryの先読み
xyuv = []
colors = [[] for _ in range(N)]
for _ in range(Q) : 
    x, y, u, v = map(int, input().split())
    x, u, v = x-1, u-1, v-1
    xyuv.append([x, y, u, v])
    colors[u].append(x)
    colors[v].append(x)
    colors[lca(u, v)].append(x)

for i in range(N) :
    colors[i] = list(set(colors[i]))

mem = dict()
C = [0] * N
D = [0] * N

def dfs(cur, pre) :
    for color in colors[cur] :
        mem[(cur, color)] = [C[color], D[color]]
    
    for nex, c, d in path[cur] :
        if nex == pre :
            continue
        C[c] += 1
        D[c] += d
        dfs(nex, cur)
        C[c] -= 1
        D[c] -= d

dfs(0, -1)
      
for x, y, u, v in xyuv :
    lca_ = lca(u, v)
    dist_ = dist[u] + dist[v] - 2*dist[lca_]
    C_ = mem[(u, x)][0] + mem[(v, x)][0] - 2*mem[(lca_, x)][0]
    D_ = mem[(u, x)][1] + mem[(v, x)][1] - 2*mem[(lca_, x)][1]
    print(dist_ - D_ + C_ * y)

Buggy submission (verdict: Runtime Error):
from collections import deque

N, Q = map(int, input().split())

path = [[] for _ in range(N)]

for _ in range(N-1) :
    a, b, c, d = (int(i) for i in input().split())
    path[a-1].append((b-1, c-1, d))
    path[b-1].append((a-1, c-1, d))

# doublingに必要なKを求める
for K in range(18) :
    if 2 ** K >= N :
        break

# bfs
parent = [[-1] * N for i in range(K)]
rank = [-1 for _ in range(N)]
dist = [0 for _ in range(N)]

rank[0] = 0
queue = deque([0])

while queue :
    cur = queue.popleft()
    for nex, _, d in path[cur] :
        if rank[nex] < 0 :
            queue.append(nex)
            parent[0][nex] = cur
            rank[nex] = rank[cur] + 1
            dist[nex] = dist[cur] + d

# doubling        
for i in range(1, K) :
    for j in range(N) :
        if parent[i-1][j] > 0 :
            parent[i][j] = parent[i-1][parent[i-1][j]]

# lca
def lca(a, b) :
    if rank[a] > rank[b] :
        a, b = b, a

    diff = rank[b] - rank[a]
    i = 0
    while diff > 0 :
        if diff & 1 :
            b = parent[i][b]
        diff >>= 1
        i += 1
        
    if a == b :
        return a

    for i in range(K-1, -1, -1) :
        if parent[i][a] != parent[i][b] :
            a = parent[i][a]
            b = parent[i][b]

    return parent[0][a]

# Queryの先読み
xyuv = []
colors = [[] for _ in range(N)]
for _ in range(Q) : 
    x, y, u, v = map(int, input().split())
    x, u, v = x-1, u-1, v-1
    xyuv.append([x, y, u, v])
    colors[u].append(x)
    colors[v].append(x)
    colors[lca(u, v)].append(x)

for i in range(N) :
    colors[i] = list(set(colors[i]))

mem = dict()
C = [0] * (N-1)
D = [0] * (N-1)

def dfs(cur, pre) :
    for color in colors[cur] :
        mem[(cur, color)] = [C[color], D[color]]
    
    for nex, c, d in path[cur] :
        if nex == pre :
            continue
        C[c] += 1
        D[c] += d
        dfs(nex, cur)
        C[c] -= 1
        D[c] -= d

dfs(0, -1)
      
for x, y, u, v in xyuv :
    lca_ = lca(u, v)
    dist_ = dist[u] + dist[v] - 2*dist[lca_]
    C_ = mem[(u, x)][0] + mem[(v, x)][0] - 2*mem[(lca_, x)][0]
    D_ = mem[(u, x)][1] + mem[(v, x)][1] - 2*mem[(lca_, x)][1]
    print(dist_ - D_ + C_ * y)

Buggy submission (verdict: Runtime Error):
from collections import deque

N, Q = map(int, input().split())

for K in range(18) :
    if 2 ** K >= N :
        break

path = [[] for _ in range(N)]

for _ in range(N-1) :
    a, b, c, d = (int(i) for i in input().split())
    path[a-1].append((b-1, c-1, d))
    path[b-1].append((a-1, c-1, d))

# Queryを先読みして必要な色を求める
xyuv = []
colors = [-1] * (N-1)
for _ in range(Q) : 
    x, y, u, v = map(int, input().split())
    xyuv.append([x-1, y, u-1, v-1])
    colors[x-1] = 1
  
count = 0
for i in range(N-1) :
    if colors[i] >= 0 :
        colors[i] = count
        count += 1

# dfs
parent = [[-1] * N for i in range(K)]
rank = [-1 for _ in range(N)]
dist = [0 for _ in range(N)]
C = [None for _ in range(N)]
D = [None for _ in range(N)]

rank[0] = 0
C[0] = [0] * count
D[0] = [0] * count
queue = deque([0])

while queue :
    cur = queue.popleft()
    
    for nex, c, d in path[cur] :
        if rank[nex] < 0 :
            parent[0][nex] = cur
            queue.append(nex)
            
            rank[nex] = rank[cur] + 1
            dist[nex] = dist[cur] + d
            
            if colors[c] >= 0 :
                C[nex] = C[cur][:]
                D[nex] = D[cur][:]
                C[nex][colors[c]] += 1
                D[nex][colors[c]] += d
            else :
                C[nex] = C[cur]
                D[nex] = D[cur]

# doubling
for i in range(1, K) :
    for j in range(N) :
        if parent[i-1][j] > 0 :
            parent[i][j] = parent[i-1][parent[i-1][j]]

def lca(a, b) :
    if rank[a] > rank[b] :
        a, b = b, a

    diff = rank[b] - rank[a]
    i = 0
    while diff > 0 :
        if diff & 1 :
            b = parent[i][b]
        diff >>= 1
        i += 1
        
    if a == b :
        return a

    for i in range(K-1, -1, -1) :
        if parent[i][a] != parent[i][b] :
            a = parent[i][a]
            b = parent[i][b]

    return parent[0][a]
    
for x, y, u, v in xyuv :
    lca_ = lca(u, v)
    color = colors[x]
    
    dist_ = dist[u] + dist[v] - 2 * dist[lca_]
    C_ = C[u][color] + C[v][color] - 2 * C[lca_][color]
    D_ = D[u][color] + D[v][color] - 2 * D[lca_][color]
    
    print(dist_ - D_ + C_ * y)

Buggy submission (verdict: Runtime Error):
from collections import deque, defaultdict

N, Q = map(int, input().split())

for K in range(18) :
    if 2 ** K >= N :
        break

path = [[] for _ in range(N)]

for _ in range(N-1) :
    a, b, c, d = (int(i) for i in input().split())
    path[a-1].append((b-1, c-1, d))
    path[b-1].append((a-1, c-1, d))

parent = [[-1] * N for i in range(K)]
rank = [(-1, [0] * N, 0) for _ in range(N)]
rank[0] = [0, [0] * N, [0] * N]

queue = deque([0])

while queue :
    cur = queue.popleft()
    r, cl, dl = rank[cur]
    
    for nex, c, d in path[cur] :
        if rank[nex][0] < 0 :
            parent[0][nex] = cur
            queue.append(nex)
            cl_ = cl[:]
            dl_ = dl[:]
            cl_[c] += 1
            dl_[c] += d
            rank[nex] = [r + 1, cl_, dl_]

# doubling
for i in range(1, K) :
    for j in range(N) :
        if parent[i-1][j] > 0 :
            parent[i][j] = parent[i-1][parent[i-1][j]]

def lca(a, b) :
    if rank[a][0] > rank[b][0] :
        a, b = b, a

    diff = rank[b][0] - rank[a][0]
    i = 0
    while diff > 0 :
        if diff & 1 :
            b = parent[i][b]
        diff >>= 1
        i += 1
        
    if a == b :
        return a

    for i in range(K-1, -1, -1) :
        if parent[i][a] != parent[i][b] :
            a = parent[i][a]
            b = parent[i][b]

    return parent[0][a]

for _ in range(Q) : 
    x, y, u, v = map(int, input().split())
    x, u, v = x-1, u-1, v-1
    cl = [0] * N
    dl = [0] * N
    for i in range(N) :        
        dl[i] += rank[u][2][i]
        dl[i] += rank[v][2][i]
        dl[i] -= rank[lca(u, v)][2][i]
    dl[x] = y * (rank[u][1][x] + rank[v][1][x] - 2 * rank[lca(u, v)][1][x])
    print(sum(dl))