#1191446 (PyPy3) No.3764 Graduation Live

• 問題: No.3764 Graduation Live
• ユーザー: ei1333333
• 提出日時: 2026-10-04 19:40:57
• 言語: PyPy3 (7.3.23 + ACL)
• 結果: AC
• 実行時間: 1,019 ms / 2,000 ms
• コード長: 3,863 bytes
• コンパイル時間: 260 ms
• コンパイル使用メモリ: 83,688 KB
• 実行使用メモリ: 174,776 KB
• 最終ジャッジ日時: 2026-10-09 20:53:33
• 合計ジャッジ時間: 37,939 ms
• ジャッジサーバーID （参考情報）: judge4_0 / judge3_0

ファイルパターン	結果
sample	AC * 3
other	AC * 45

ソースコード

import sys

def solve_case(n, limits, values):
    # 0: 始点、1..4: 左側、5..8: 右側、9: 終点
    source, sink = 0, 9
    vertex_count = 10
    adj = [[] for _ in range(vertex_count)]
    to = []
    cap = []
    cost = []
    edge_values = []

    def add_edge(u, v, capacity, songs=None):
        e = len(to)
        adj[u].append(e)
        adj[v].append(e ^ 1)
        to.extend((v, u))
        cap.extend((capacity, 0))
        cost.extend((-songs[0] if songs else 0, 0))
        edge_values.extend((songs, songs))

    for i, limit in enumerate(limits):
        add_edge(source, 1 + i, limit)
        add_edge(5 + i, sink, limit)

    endpoints = ((0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3))
    for songs, (u, v) in zip(values, endpoints):
        if not songs:
            continue
        songs.sort(reverse=True)
        m = len(songs)
        add_edge(1 + u, 5 + v, m, songs)
        add_edge(1 + v, 5 + u, m, songs)

    # 初期状態は頂点番号順がトポロジカル順になっている DAG。
    # 負辺に対応するため、最初のポテンシャルを求める。
    potential = [0] * vertex_count
    for u in range(vertex_count):
        for e in adj[u]:
            if cap[e]:
                v = to[e]
                nd = potential[u] + cost[e]
                if nd < potential[v]:
                    potential[v] = nd

    inf = 10**30
    total_cost = 0
    flow = 0
    answers = []
    prev = [-1] * vertex_count
    search_order = range(sink - 1, -1, -1)

    while flow < 2 * n:
        dist = [inf] * vertex_count
        used = [False] * vertex_count
        dist[source] = 0

        # 頂点数が 10 なので、ヒープを使わず最小距離の頂点を探す。
        while True:
            # 同距離なら終点を優先する。
            u = sink
            best = dist[sink]
            for v in search_order:
                if not used[v] and dist[v] < best:
                    u = v
                    best = dist[v]

            if u == sink:
                break

            used[u] = True
            base = best + potential[u]
            for e in adj[u]:
                if cap[e]:
                    v = to[e]
                    nd = base + cost[e] - potential[v]
                    if nd < dist[v]:
                        dist[v] = nd
                        prev[v] = e

        if dist[sink] == inf:
            break

        # 終点で探索を打ち切っているので、更新量を dist[sink] 以下にする。
        d = dist[sink]
        for v in range(vertex_count):
            potential[v] += dist[v] if dist[v] < d else d

        total_cost += potential[sink] - potential[source]

        v = sink
        while v != source:
            e = prev[v]
            cap[e] -= 1
            cap[e ^ 1] += 1

            songs = edge_values[e]
            if songs is not None:
                fwd = e & ~1
                k = len(songs) - cap[fwd]
                if cap[fwd]:
                    cost[fwd] = -songs[k]
                if cap[fwd ^ 1]:
                    cost[fwd ^ 1] = songs[k - 1]

            v = to[e ^ 1]

        flow += 1
        if flow % 2 == 0:
            answers.append((-total_cost) // 2)

    return answers


def main():
    data = iter(map(int, sys.stdin.buffer.read().split()))
    test_count = next(data)
    output = []

    for _ in range(test_count):
        n = next(data)
        limits = [n, next(data), next(data), next(data)]
        values = [[] for _ in range(6)]
        for _ in range(n):
            t = next(data)
            v = next(data)
            values[t - 1].append(v)
        answers = solve_case(n, limits, values)
        output.append(' '.join(map(str, [len(answers)] + answers)))

    sys.stdout.write('\n'.join(output))
    sys.stdout.write('\n')


if __name__ == '__main__':
    main()