Problem p02986 (python)

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

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)]

def set_ancestor(x=1, d=0, parent=0):
    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
        set_ancestor(y, d+1, x)

depth = [0] * (N+1)
U = 17
ancestor = [[0] * (U+1) for _ in range(N+1)] # 2**i だけ遡った頂点の番号、行き過ぎは0
set_ancestor()

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, sums = defaultdict(int), nums = defaultdict(int), 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(1)

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

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)
import numpy as np

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

def dfs(x, tour_v, tour_d, depth = 0, parent=None):
    tour_v.append(x)
    tour_d.append(depth)
    for y in graph[x]:
        if y == parent:
            continue
        dfs(y, tour_v, tour_d, depth + 1, parent = x)
        tour_v.append(x)
        tour_d.append(depth)

tour_v = []
tour_d = []
dfs(1, tour_v, tour_d)

v_to_i = {v:i for i,v in enumerate(tour_v)}
v_to_d = dict(zip(tour_v, tour_d))
tour_v = np.array(tour_v)
tour_d = np.array(tour_d)

# 各インデックスから、2^i 項目分見て、一番depthの小さいインデックス
L = len(tour_v)
U = L.bit_length()
sparse_table = [np.arange(L)]
for i in range(1,U):
    prev = sparse_table[-1]
    width = 1 << (i-1)
    n1 = prev[:-width]
    n2 = prev[width:]
    sparse_table.append(np.where(tour_d[n1] < tour_d[n2], n1, n2))

def LCA(x,y):
    ix = v_to_i[x]
    iy = v_to_i[y]
    if ix > iy:
        ix,iy = iy,ix
    items = iy - ix + 1
    L = items.bit_length() - 1
    n1 = sparse_table[L][ix]
    n2 = sparse_table[L][iy - (1 << L) + 1]
    n = n1 if tour_d[n1] < tour_d[n2] else n2
    return tour_v[n]

answer = [v_to_d[x] + v_to_d[y] - 2 * v_to_d[LCA(x,y)] + 1 for x,y in query]
print('\n'.join(map(str,answer)))

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)
from collections import defaultdict

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)]

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

depth = [0] * (N+1)
ancestor = [[] for _ in range(N+1)] # 2**i だけ遡った頂点の番号、行き過ぎは0
set_ancestor()

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, sums = defaultdict(int), nums = defaultdict(int), 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(1)

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

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)
from collections import defaultdict
from functools import lru_cache

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)]

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

depth = [0] * (N+1)
ancestor = [[] for _ in range(N+1)] # 2**i だけ遡った頂点の番号、行き過ぎは0
set_ancestor()

@lru_cache(None)
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,-1))

answer = [0] * Q# 頂点、色に対して、距離、自色距離、自色個数)を入れる

def dfs(x, sums = defaultdict(int), nums = defaultdict(int), 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(1)

print('\n'.join(str(x) for x in answer))