Problem p02986 (python)

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

N, Q = map(int, input().split())
graph = [[] for _ in range(N+1)]
for _ in range(N-1):
    a, b, c, d = map(int, input().split())
    graph[a].append((b, c, d))
    graph[b].append((a, c, d))
query = [[int(x) for x in input().split()] for _ in range(Q)]

depth = [0] * (N+1)
U = 17
ancestor = [[0] * (U+1) for _ in range(N+1)]

q = [(1, 0, 0)]
while q:
    qq = []
    for x, d, parent in q:
        depth[x] = d
        ax = ancestor[x]
        ax[0] = parent
        for i in range(d.bit_length()-1):
            ax[i+1] = ancestor[ax[i]][i]
        for y, _, _ in graph[x]:
            if y == parent:
                continue
            qq.append((y, d+1, x))
    q = qq

def LCA(x, y):
    dx = depth[x]
    dy = depth[y]
    if dx > dy:
        x, y = y, x
        dx, dy = dy, dx
    diff = dy - dx
    while diff:
        step = diff & (-diff)
        y = ancestor[y][step.bit_length()-1]
        diff -= step
    while x != y:
        j = 0
        while ancestor[x][j] != ancestor[y][j]:
            j += 1
        if j == 0:
            return ancestor[x][0]
        x = ancestor[x][j-1]
        y = ancestor[y][j-1]
    return x

tasks = [[] for _ in range(N+1)]
for i, (x, y, u, v) in enumerate(query):
    tasks[u].append((i, x, y, 1))
    tasks[v].append((i, x, y, 1))
    tasks[LCA(u, v)].append((i, x, y, -2))

answer = [0] * Q

def dfs(x = 1, sums = [0] * (N+1), nums = [0] * (N+1), total = 0, parent = 0):
    for i, c, d, coef in tasks[x]:
        answer[i] += coef * (total - sums[c] + nums[c] * d)
    for y, c, d in graph[x]:
        if y == parent:
            continue
        sums[c] += d
        nums[c] += 1
        total += d
        dfs(y, sums, nums, total, x)
        sums[c] -= d
        nums[c] -= 1
        total -= d
    return

dfs()

print('\n'.join(map(str, answer)))

Buggy submission (verdict: Runtime Error):
import sys

N, Q = map(int, input().split())
edge = [[] for _ in range(N+1)]
for _ in range(N-1):
        a, b, c, d = map(int, input().split())
        edge[a].append((b, c, d))
        edge[b].append((a, c, d))

depth = [0] * (N+1)
U = 17
ancestor = [[0] * (U+1) for _ in range(N+1)]

q = [(1, 0, 0)]
while q:
        qq = []
        for x, d, parent in q:
                depth[x] = d
                ax = ancestor[x]
                ax[0] = parent
                for i in range(d.bit_length()-1):
                        ax[i+1] = ancestor[ax[i]][i]

                for y, _, _ in edge[x]:
                        if y == parent:
                                continue
                        qq.append((y, d+1, x))
        q = qq

def lca(x, y):
        dx = depth[x]
        dy = depth[y]
        if dx > dy:
                x, y = y, x
                dx, dy = dy, dx
        diff = dy - dx
        while diff:
                step = diff & (-diff)
                y = ancestor[y][step.bit_length()-1]
                diff -= step
        while x!=y:
                j = 0
                while ancestor[x][j] != ancestor[y][j]:
                        j += 1
                if j==0:
                        return ancestor[x][0]
                x = ancestor[x][j-1]
                y = ancestor[y][j-1]
        return x

task = [[] for _ in range(N+1)]
for i in range(Q):
        x, y, u, v = map(int, input().split())
        task[u].append((i, x, y, 1))
        task[v].append((i, x, y, 1))
        task[lca(u, v)].append((i, x, y, -2))

answer = [0] * Q

def dfs(x=1, sums=[0]*(N+1), nums=[0]*(N+1), total=0, parent=0):
        for i, c, d, coef in task[x]:
                answer[i] += coef * (total - sums[c] + nums[c]*d)
        for y, c, d in edge[x]:
                if y == parent:
                        continue
                sums[c] += d
                nums[c] += 1
                total += d
                dfs(y, sums, nums, total, x)
                sums[c] -= d
                nums[c] -= d
                total -= d
        return

dfs()

print('\n'.join(map(str, answer)))

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

N, Q = map(int, input().split())
edge = defaultdict(list)
for i in range(N-1):
        a, b, c, d = map(int, input().split())
        edge[a].append((b, c, d))
        edge[b].append((a, c, d))

depth = [0] * (N+1)
dist_sum = [0] * (N+1)
parent = list(range(N+1))

#dfs
stack = []
stack.append((1, 1, 0, 0))
while stack:
        now, prev, dep, dist = stack.pop()
        depth[now] = dep
        parent[now] = prev
        dist_sum[now] = dist
        for e, c, d in edge[now]:
                if e == prev:
                        continue
                else:
                        stack.append((e, now, dep+1, dist+d))

U = 17
doubling = [[1] * U for i in range(N+1)]
for i in range(U):
        for j in range(1, N+1):
                doubling[j][i] = parent[j]
        parent = [parent[parent[i]] for i in range(N+1)]

def lca(a, b):
        x, y = depth[a], depth[b]
        if x > y:
                x, y = y, x
                a, b = b, a        
        diff = y-x
        if diff > 0:
                for i in reversed(range(U)):
                        if diff > 2**i:
                                b = doubling[b][i]
                                diff -= 2**i
                b = doubling[b][0]
        if a == b:
                return a
        
        for i in reversed(range(int(math.log2(x)) + 2)):
                if doubling[a][i] != doubling[b][i]:
                        a, b = doubling[a][i], doubling[b][i]
        return doubling[a][0]

task = defaultdict(list)
for i in range(Q):
        x, y, u, v = map(int, input().split())
        task[u].append((i, x, y, 1))
        task[v].append((i, x, y, 1))
        task[lca(u, v)].append((i, x, y, -2))

stack = []
ans = [0] * Q
color_num = [0] * N
color_dist = [0] * N
stack.append((1, -1, color_num, color_dist))
while stack:
        now, prev, color_num, color_dist = stack.pop()
        if len(task[now]) > 0:
                for i, x, y, coeff in task[now]:
                        ans[i] += coeff * (dist_sum[now] - color_dist[x] + color_num[x] * y)
                        
        for e, c, d in edge[now]:
                if e == prev:
                        continue
                color_num[c] += 1
                color_dist[c] += d
                stack.append((e, now, color_num[:], color_dist[:]))
                color_num[c] -= 1
                color_dist[c] -= d

print('\n'.join([str(ans[i]) for i in range(Q)]))

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

N, Q = map(int, input().split())
edge = defaultdict(list)
for i in range(N-1):
        a, b, c, d = map(int, input().split())
        edge[a].append((b, c, d))
        edge[b].append((a, c, d))

depth = [0] * (N+1)
dist_sum = [0] * (N+1)
parent = list(range(N+1))

#dfs
stack = []
stack.append((1, -1, 0, 0))
while stack:
        now, prev, dep, dist = stack.pop()
        depth[now] = dep
        parent[now] = prev
        dist_sum[now] = dist
        for e, c, d in edge[now]:
                if e == prev:
                        continue
                else:
                        stack.append((e, now, dep+1, dist+d))

doubling = [[1] * N for i in range(N+1)]
for i in range(N):
        for j in range(1, N+1):
                doubling[j][i] = parent[j]
        parent = [parent[parent[i]] for i in range(N+1)]

def lca(a, b):
        x, y = depth[a], depth[b]
        if x > y:
                x, y = y, x
                a, b = b, a        
        diff = y-x
        if diff > 0:
                for i in reversed(range(N)):
                        if diff > 2**i:
                                b = doubling[b][i]
                                diff -= 2**i
                b = doubling[b][0]
        if a == b:
                return a
        
        for i in reversed(range(int(math.log2(x)) + 2)):
                if doubling[a][i] != doubling[b][i]:
                        a, b = doubling[a][i], doubling[b][i]
        return doubling[a][0]

task = defaultdict(list)
for i in range(Q):
        x, y, u, v = map(int, input().split())
        task[u].append((i, x, y, 1))
        task[v].append((i, x, y, 1))
        task[lca(u, v)].append((i, x, y, -2))

stack = []
ans = [0] * Q
color_num = [0] * N
color_dist = [0] * N
stack.append((1, -1, color_num, color_dist))
while stack:
        now, prev, color_num, color_dist = stack.pop()
        if len(task[now]) > 0:
                for i, x, y, coeff in task[now]:
                        ans[i] += coeff * (dist_sum[now] - color_dist[x] + color_num[x] * y)
                        
        for e, c, d in edge[now]:
                if e == prev:
                        continue
                color_num[c] += 1
                color_dist[c] += d
                stack.append((e, now, color_num[:], color_dist[:]))
                color_num[c] -= 1
                color_dist[c] -= d

print('\n'.join([str(ans[i]) for i in range(Q)]))

Buggy submission (verdict: Runtime Error):
from collections import defaultdict, deque
from math import *
import sys
sys.setrecursionlimit(10**8)

N, Q = map(int, input().split())
edge = defaultdict(list)
for i in range(N-1):
    a, b, c, d = map(int, input().split())
    edge[a].append((b, c, d))
    edge[b].append((a, c, d))

query = []
for i in range(Q):
    x, y, u, v = map(int, input().split())
    query.append((x, y, u, v))

def dfs(prev, cur, cur_color_num, cur_color_total, cur_depth):
    parent[cur] = prev
    depth[cur] = cur_depth
    color_num[cur] = cur_color_num[:]
    color_total[cur] = cur_color_total[:]
    for e, c, d in edge[cur]:
        if e != prev:
            temp_num = cur_color_num[:]
            temp_total = cur_color_total[:]
            temp_num[c] += 1
            temp_total[c] += d
            dist[e] = dist[cur] + d
            dfs(cur, e, temp_num, temp_total, cur_depth+1)
    return

color_num =[[0] * N for i in range(N+1)]
color_total = [[0] * N for i in range(N+1)]
depth = [0] * (N+1)
dist = [0] * (N+1)
parent = list(range(N+1))
cur_color_num = [0] * N
cur_color_total = [0] * N
dfs(1, 1, cur_color_num, cur_color_total, 0)

doubling = [parent]
k = int(log2(N)) + 1
for i in range(k):
    parent = [parent[parent[i]] for i in range(N+1)]
    doubling.append(parent[:])

for i in range(Q):
    x, y, u, v = query[i]
    cur_u = u
    cur_v = v

    diff = depth[u] - depth[v]
    if diff > 0:
        for j in range(k):
            ind = k-j-1
            if 2**ind <= diff:
                cur_u = doubling[ind][cur_u]
                diff -= 2**ind
    else:
        diff = -diff
        for j in range(k):
            ind = k-j-1
            if 2**ind <= diff:
                cur_v = doubling[ind][cur_v]
                diff -= 2**ind

    for j in range(k):
        ind = k-j-1
        if doubling[ind][cur_u] != doubling[ind][cur_v]:
            cur_u = doubling[ind][cur_u]
            cur_v = doubling[ind][cur_v]

    lca = doubling[0][cur_u]
    dist_lca = dist[lca] - color_total[lca][x] + color_num[lca][x] * y
    dist_u = dist[u] - color_total[u][x] + color_num[u][x] * y
    dist_v = dist[v] - color_total[v][x] + color_num[v][x] * y
    #print(str(dist_u) + ' ' + str(dist_v) + ' '  +  str(dist_lca))
    #print(str(dist[v]) + ' ' + str(color_total[v][x]) + ' '  +  str(color_num[v][x]))
    print(dist_u + dist_v - dist_lca)

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

N, Q = map(int, input().split())
edge = defaultdict(list)
for i in range(N-1):
    a, b, c, d = map(int, input().split())
    edge[a].append((b, c, d))
    edge[b].append((a, c, d))

query = []
for i in range(Q):
    x, y, u, v = map(int, input().split())
    query.append((x, y, u, v))

def dfs(prev, cur, cur_color_num, cur_color_total, cur_depth):
    parent[cur] = prev
    depth[cur] = cur_depth
    color_num[cur] = cur_color_num
    color_total[cur] = cur_color_total
    for e, c, d in edge[cur]:
        if e != prev:
            temp_num = cur_color_num[:]
            temp_total = cur_color_total[:]
            temp_num[c] += 1
            temp_total[c] += d
            dist[e] = dist[cur] + d
            dfs(cur, e, temp_num, temp_total, cur_depth+1)
    return

color_num =[[0] * N for i in range(N+1)]
color_total = [[0] * N for i in range(N+1)]
depth = [0] * (N+1)
dist = [0] * (N+1)
parent = list(range(N+1))
cur_color_num = [0] * N
cur_color_total = [0] * N
dfs(1, 1, cur_color_num, cur_color_total, 0)

doubling = [parent]
k = int(log2(N)) + 1
for i in range(k):
    parent = [parent[parent[i]] for i in range(N+1)]
    doubling.append(parent)

for i in range(Q):
    x, y, u, v = query[i]
    cur_u = u
    cur_v = v

    diff = depth[u] - depth[v]
    if diff > 0:
        for j in range(k):
            ind = k-j-1
            if 2**ind <= diff:
                cur_u = doubling[ind][cur_u]
                diff -= 2**ind
    else:
        diff = -diff
        for j in range(k):
            ind = k-j-1
            if 2**ind <= diff:
                cur_v = doubling[ind][cur_v]
                diff -= 2**ind

    for j in range(k):
        ind = k-j-1
        if doubling[ind][cur_u] != doubling[ind][cur_v]:
            cur_u = doubling[ind][cur_u]
            cur_v = doubling[ind][cur_v]

    lca = doubling[0][cur_u]
    dist_lca = dist[lca] - color_total[lca][x] + color_num[lca][x] * y
    dist_u = dist[u] - color_total[u][x] + color_num[u][x] * y
    dist_v = dist[v] - color_total[v][x] + color_num[v][x] * y
    #print(str(dist_u) + ' ' + str(dist_v) + ' '  +  str(dist_lca))
    #print(str(dist[v]) + ' ' + str(color_total[v][x]) + ' '  +  str(color_num[v][x]))
    print(dist_u + dist_v - dist_lca)