🎇

ABC267 E - Erasing Vertices 2 Python解答例

2022/09/04に公開約2,800字

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

実際の提出はこちら

GitHubで編集を提案

Discussion

ログインするとコメントできます