Problem p03020 (python)

Fixed submission:
import sys
from collections import deque


def min_cost_flow(links, links_from, s, t, flow, n2):
    remain = flow
    result = 0
    INF = 10 ** 12
    predecessors = [0] * n2  # link_id

    while remain:
        # print(remain)
        distances = [INF] * n2
        updated = [False] * n2
        distances[s] = 0
        updated[s] = True
        q = deque([s])
        while q:
            v = q.popleft()
            vc = distances[v]
            updated[v] = False

            for li in links_from[v]:
                _, u, cap, cost = links[li]
                if cap == 0:
                    continue
                new_cost = vc + cost
                if new_cost >= distances[u]:
                    continue
                distances[u] = new_cost
                predecessors[u] = li
                if not updated[u]:
                    updated[u] = True
                    q.append(u)

        min_cap = remain
        v = t
        while v != s:
            li = predecessors[v]
            l = links[li]
            min_cap = min(min_cap, l[2])
            v = l[0]
        v = t
        while v != s:
            li = predecessors[v]
            l = links[li]
            l[2] -= min_cap
            links[li ^ 1][2] += min_cap
            v = l[0]
        remain -= min_cap
        result -= min_cap * distances[t]
    return result


n = int(input())
lines = sys.stdin.readlines()
n2 = 2 * n + 6
s, t = 0, 2 * n + 1
k1, k2, k3, k4 = range(n2 - 4, n2)
balls = 0
links_from = [[] for _ in range(n2)]
links = []  # [[src, tgt, capacity, unit_cost], ]


def add_link(s, t, cap, cost):
    i = len(links)
    links.append([s, t, cap, cost])
    links.append([t, s, 0, -cost])
    links_from[s].append(i)
    links_from[t].append(i + 1)


for i in range(n):
    ri = i + 1
    rx, ry, rc = map(int, lines[i].split())
    balls += rc
    add_link(s, ri, rc, 0)
    add_link(ri, k1, rc, rx + ry)
    add_link(ri, k2, rc, -rx + ry)
    add_link(ri, k3, rc, rx - ry)
    add_link(ri, k4, rc, -rx - ry)

for i in range(n, 2 * n):
    bi = i + 1
    bx, by, bc = map(int, lines[i].split())
    add_link(bi, t, bc, 0)
    add_link(k1, bi, bc, -bx - by)
    add_link(k2, bi, bc, bx - by)
    add_link(k3, bi, bc, -bx + by)
    add_link(k4, bi, bc, bx + by)

print(min_cost_flow(links, links_from, s, t, balls, n2))

Buggy submission (verdict: Runtime Error):
import sys
from heapq import heappop, heappush


def min_cost_flow(links, s, t, flow, n2, INF, REV):
    remain = flow
    ans = 0

    potentials = [0] * n2  # to adjust negative edge for dijkstra
    predecessors_v = [0] * n2
    predecessors_i = [0] * n2

    while remain:
        distances = [INF] * n2
        distances[s] = 0
        q = [(0, s)]  # (cost, v, (p, v_idx_in_p))
        while q:
            cost, v = heappop(q)
            if cost > distances[v]:
                continue

            if v == t:
                break

            for i, (u, cap, uc, j) in enumerate(links[v]):
                if cap == 0:
                    continue
                new_cost = cost + uc + potentials[v] - potentials[u]
                if distances[u] <= new_cost:
                    continue
                distances[u] = new_cost
                predecessors_v[u] = v
                predecessors_i[u] = i
                heappush(q, (new_cost, u))

        # if distances[t] == INF:
        #     return -1

        for i in range(n2):
            potentials[i] += distances[i]

        min_cap = remain
        v = t
        forward_links = []
        backward_links = []
        while v != s:
            p = predecessors_v[v]
            i = predecessors_i[v]
            l = links[p][i]
            forward_links.append(l)
            if v != t:
                backward_links.append(links[v][l[3]])
            min_cap = min(min_cap, l[1])
            v = p
        for l in forward_links:
            l[1] -= min_cap
        for l in backward_links:
            l[1] += min_cap
        remain -= min_cap
        ans += min_cap * (REV * 2 - potentials[t])
    return ans


n = int(input())
lines = sys.stdin.readlines()
flow = 0
INF = 10 ** 18
REV = 10 ** 9
n2 = 2 * n + 6
t = 2 * n + 1
links = [[] for _ in range(n2)]  # [[tgt_node, capacity, unit_cost, idx_in_tgt], ]
for ri, line in enumerate(lines[:n], start=1):
    rx, ry, rc = map(int, line.split())
    flow += rc
    links[0].append([ri, rc, 0, 0])
    ki = n2 - 4
    kc = REV - (rx + ry)
    links[ri].append([ki, INF, kc, len(links[ki])])
    links[ki].append([ri, 0, -kc, len(links[ri]) - 1])
    ki = n2 - 3
    kc = REV - (-rx + ry)
    links[ri].append([ki, INF, kc, len(links[ki])])
    links[ki].append([ri, 0, -kc, len(links[ri]) - 1])
    ki = n2 - 2
    kc = REV - (rx - ry)
    links[ri].append([ki, INF, kc, len(links[ki])])
    links[ki].append([ri, 0, -kc, len(links[ri]) - 1])
    ki = n2 - 1
    kc = REV - (-rx - ry)
    links[ri].append([ki, INF, kc, len(links[ki])])
    links[ki].append([ri, 0, -kc, len(links[ri]) - 1])

for bi, line in enumerate(lines[n:], start=n + 1):
    bx, by, bc = map(int, line.split())
    links[bi].append([t, bc, 0, 0])
    ki = n2 - 4
    kc = REV - (-bx - by)
    links[ki].append([bi, INF, kc, len(links[bi])])
    links[bi].append([ki, 0, -kc, len(links[ki]) - 1])
    ki = n2 - 3
    kc = REV - (bx - by)
    links[ki].append([bi, INF, kc, len(links[bi])])
    links[bi].append([ki, 0, -kc, len(links[ki]) - 1])
    ki = n2 - 2
    kc = REV - (-bx + by)
    links[ki].append([bi, INF, kc, len(links[bi])])
    links[bi].append([ki, 0, -kc, len(links[ki]) - 1])
    ki = n2 - 1
    kc = REV - (bx + by)
    links[ki].append([bi, INF, kc, len(links[bi])])
    links[bi].append([ki, 0, -kc, len(links[ki]) - 1])

print(min_cost_flow(links, 0, t, flow, n2, INF, REV))