ABC267 E - Erasing Vertices 2 Python解答例

AtCoder Beginner Contest E - Erasing Vertices 2をPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

$N$頂点$M$辺の単純無向グラフ（すなわち、自己辺も多重辺もない無向グラフ）が与えられます。$i$本目の辺は頂点$U_i$と頂点$V_i$を結んでいます。頂点$i$には正整数$A_i$が書かれています。

あなたは、以下の操作を$N$回繰り返します。

• まだ削除されていない頂点$x$を選び、「頂点$x$」と「頂点$x$を端点に持つ辺全て」を削除する。この操作のコストは、頂点$x$と辺で直接結ばれていて、かつまだ削除されていない頂点に書かれている整数の総和である。

$N$回の操作全体のコストを、$1$回ごとの操作におけるコストのうちの最大値として定めます。操作全体のコストとして取り得る値の最小値を求めてください。

制約

• $1 \le N \le 2 \times 10^5$
• $0 \le M \le 2 \times 10^5$
• $1 \le A_i \le 10^9$
• $1 \le U_i,V_i \le N$
• 与えられるグラフは単純。
• 入力は全て整数。

解答例

コストが小さい点から貪欲に

ユーザ解説に書かれている通り、どのノードを消したとしても、他のノードの削除コストが増えることはないため、「現時点でコストが最も低いノードを削除する」を$N$回繰り返すのが最適です。

まず最初の段階で各ノードの削除コストを計算しておきます。
この削除コストを、ノード番号と一緒に優先度付きキューに格納します。これにより、「現時点でコストが最も低いノード」を取得することができます。

また、ノード$i$を削除すると、隣接する頂点の削除コストが$A_i$だけ減少します。
次の操作でコストが最も低いノードを探すためには、これを反映させなければなりません。優先度付きキューの中にある要素を書き換えることは難しいので、値を更新した後のノードを新たに追加してしまいます(ダイクストラ法で使われるテクニックと同じ)。

ただし、このようにすると同じノードに関する要素が複数回キューに入ることになるので、ノードが削除されたかどうかを記録しておき、既に削除されているならスキップするようにします。

操作を繰り返していき、削除コストの最大値が答えとなります。

実装例

Pythonによる実装例を以下に示します。

e.py

from typing import *
import heapq


def main():
    # 入力受け取り
    N, M = map(int, input().split())
    A: List[int] = list(map(int, input().split()))
    # ノード番号と頂点の重みを対応付ける辞書(最初のコスト計算に使う)
    W: Dict[int, int] = {i: a for i, a in enumerate(A)}
    # グラフの構築
    graph: List[Set[int]] = [set() for i in range(N)]
    for i in range(M):
        u, v = map(int, input().split())
        u -= 1
        v -= 1
        graph[u].add(v)
        graph[v].add(u)

    # 初期状態における、各ノードの削除コスト
    # 隣接するノードの重みの総和(sum[W[node] for node in graph[i]]の部分)を全てのノードについて計算する
    cost: List[int] = [sum([W[node] for node in graph[i]]) for i in range(N)]
    # ノードが削除されたかどうかを記録する配列
    eliminated: List[bool] = [False] * N

    # 優先度付きキュー。削除コストと頂点番号を一緒に格納する
    candidate: List[Tuple[int, int]] = [(c, i) for i, c in enumerate(cost)]
    # ヒープの構築
    heapq.heapify(candidate)

    # キューが空になるまで操作する
    answer: int = 0
    while candidate:
        # 現時点で削除コストが最小のノードを取得する
        c, node = heapq.heappop(candidate)

        # そのノードが既に削除されていたら何もせずスキップする
        if eliminated[node]:
            continue

        # 隣接する頂点の削除コストを減らす
        for next_node in graph[node]:
            # nodeを削除すると隣接する頂点の削除コストはA[node]だけ減る
            cost[next_node] -= A[node]
            # 削除コストを更新したのでキューに追加する
            heapq.heappush(candidate, (cost[next_node], next_node))

        # 最大値を記録していく
        answer = max(answer, c)
        # ノードが削除されたことを記録する
        eliminated[node] = True

    # 答えの出力
    print(answer)


if __name__ == "__main__":
    main()

実際の提出はこちら。