Problem p03021 (Python) — buggy submission, verdict: Runtime Error

import sys
from heapq import nlargest

sys.setrecursionlimit(10000)


def first_dfs(v, p, stones, dists, costs):
    stone = s[v]
    dist = 0
    cost = 0
    children_dists = []
    for u in links[v]:
        if u == p:
            continue
        us, ud, uc = first_dfs(u, v, stones, dists, costs)
        stone += us
        dist += us + ud
        cost += uc
        children_dists.append((us + ud, uc))
    # print(v, p, stone, dist, cost, children_dists)
    if len(children_dists) > 1:
        md, mc = max(children_dists)
        if dist - md < md - mc * 2:
            cost = dist - md + mc
        else:
            cost = dist // 2
        # print(v, p, md, hd, dist, cost)
    stones[v] = stone
    dists[v] = dist
    costs[v] = cost
    return stone, dist, cost


def second_dfs(v, p, pd, pc):
    children_dists = []
    dist = 0

    for u in links[v]:
        if u == p:
            dist += pd
            children_dists.append((pd, pc))
            continue
        ud = stones[u] + dists[u]
        uc = costs[u]
        dist += ud
        children_dists.append((ud, uc))
    # print(v, p, pd, pc, dist, children_dists)

    # only when source and leaf node
    if len(children_dists) == 1:
        u = next(iter(links[v]))
        return second_dfs(u, v, s[v], 0)

    (md1, mc1), (md2, mc2) = nlargest(2, children_dists)
    ret = INF

    hd, m = divmod(dist, 2)
    # print(' ', hd, m, md1, mc1, dist - md1, md1 - mc1 * 2)
    if m == 0 and dist - md1 >= md1 - mc1 * 2:
        ret = min(ret, hd)

    for u in links[v]:
        if u == p:
            continue
        us = stones[u]
        ud = dists[u]
        uc = costs[u]

        # dist=0 node cannot be a center
        if ud == 0:
            continue
        usd = us + ud
        if usd == md1 and uc == mc1:
            md, mc = md2, mc2
        else:
            md, mc = md1, mc1
        td = dist - usd
        if td - md < md - mc * 2:
            vc = td - md + mc
        else:
            vc = td // 2
        result = second_dfs(u, v, td + (all_stones - us), vc)
        ret = min(ret, result)

    return ret


n = int(input())
s = list(map(int, input()))
links = [set() for _ in [0] * n]
for line in sys.stdin:
    a, b = map(int, line.split())
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)
all_stones = sum(s)
if all_stones == 1:
    print(0)
    exit()
stones = [0] * n  # 1を根とした木の、頂点vの部分木以下（自身含む）にある石の数
dists = [0] * n  # 頂点vの部分木以下にある石 x それぞれのvとの距離 の和
costs = [0] * n  # 頂点v以下の石同士で、最大限頂点vに近づけておくためにできる操作の数
first_dfs(0, -1, stones, dists, costs)
# print(stones)
# print(dists)
# print(costs)
INF = 10**10
ret = second_dfs(0, -1, 0, 0)
print(-1 if ret == INF else ret)


Error: 0

Fixed submission:

import sys

sys.setrecursionlimit(10000)


def first_dfs(v, p, stones, dists, costs):
    stone = s[v]
    dist = 0
    cost = 0
    children_dists = []
    for u in links[v]:
        if u == p:
            continue
        us, ud, uc = first_dfs(u, v, stones, dists, costs)
        stone += us
        dist += us + ud
        cost += uc
        children_dists.append((us + ud, uc))
    # print(v, p, stone, dist, cost, children_dists)
    if len(children_dists) > 1:
        md, mc = max(children_dists)
        if dist - md < md - mc * 2:
            cost = dist - md + mc
        else:
            cost = dist // 2
        # print(v, p, md, hd, dist, cost)
    stones[v] = stone
    dists[v] = dist
    costs[v] = cost
    return stone, dist, cost


def second_dfs(v, p, pd, pc):
    children_dists = []
    dist = 0

    for u in links[v]:
        if u == p:
            dist += pd
            children_dists.append((pd, pc))
            continue
        ud = stones[u] + dists[u]
        uc = costs[u]
        dist += ud
        children_dists.append((ud, uc))
    # print(v, p, pd, pc, dist, children_dists)

    if len(children_dists) == 1:
        u = next(iter(links[v]))
        second_dfs(u, v, s[v], 0)
        return

    children_dists.sort(reverse=True)
    md1, mc1 = children_dists[0]
    md2, mc2 = children_dists[1]

    hd, m = divmod(dist, 2)
    # print(' ', hd, m, md1, mc1)
    if m == 0 and dist - md1 >= md1 - mc1 * 2:
        print(hd)
        exit()

    for u in links[v]:
        if u == p:
            continue
        us = stones[u]
        ud = dists[u]
        uc = costs[u]
        usd = us + ud
        if usd > 0:
            if usd == md1 and uc == mc1:
                md, mc = md2, mc2
            else:
                md, mc = md1, mc1
            td = dist - usd
            if td - md < md - mc * 2:
                vc = td - md + mc
            else:
                vc = td // 2
            second_dfs(u, v, td + (all_stones - us), vc)


n = int(input())
s = list(map(int, input()))
links = [set() for _ in [0] * n]
for line in sys.stdin:
    a, b = map(int, line.split())
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)
all_stones = sum(s)
if all_stones == 1:
    print(0)
    exit()
# print(links)
stones = [0] * n  # 1を根とした木の、頂点vの部分木以下（自身含む）にある石の数
dists = [0] * n  # 頂点vの部分木以下にある石 x それぞれのvとの距離 の和
costs = [0] * n  # 頂点v以下同士の石で、最大限頂点vに近づけておくためにできる操作の数
first_dfs(0, -1, stones, dists, costs)
# print(stones)
# print(dists)
# print(costs)
second_dfs(0, -1, 0, 0)
print(-1)