AtCoder ABC246G「Game on Tree 3」をUnion-Findっぽく貪欲法で解く

前置き

ABC246 G 「Game on Tree 3」を自力で解いたら、解説よりもシンプル（？）な解法になったので紹介します。

atcoder.jp

公式及びユーザー解説はこちら。

atcoder.jp

解法

公式解説には木DPとか二分探索などと書かれていますが、面倒なので貪欲に考えます。

今回は、青木くんが得点の高い頂点から順にいくつまで $0$ にできるかを考えられれば良さそうです。

入力例1を使って考えます。

f:id:dnek:20220411181409p:plain — ABC246G入力例1図

まず、高橋くんが頂点 $7$ にそのまま辿り着くと最大値の $10$ 点を取られてしまうので、これを $0$ にすることが最優先事項です。後のことを考え、頂点 $7$ を $0$ にするタイミングをなるべく遅らせてみます。すなわち、頂点 $7$ の親である頂点 $5$ から開始するターンに $0$ にします。逆に、これより前のターンに頂点 $7$ を $0$ にしても先出し不利で全く意味がないので、これが最善です。

次に $0$ にすべきは $6$ 点である頂点 $4$ か $6$ のいずれかですが、この順番は関係ないのでとりあえず頂点 $6$ から考えます。先程と同様に頂点 $5$ 開始時に $0$ にしたいところですが、このタイミングは既に頂点 $7$ に使うことが決まっているので、次点で更に親の頂点 $2$ 開始時を使うことにします。このように $0$ にしたい頂点の祖先を近い方から順に見ていって、まだ使われていないタイミングがあればそこを使うようにします。

同様に、頂点 $4$ を $0$ にするタイミングは頂点 $2$ 開始時が予約済みなので頂点 $1$ 開始時となります。

次は $5$ 点である頂点 $5$ を $0$ にしたいところですが、祖先である頂点 $2$ と $1$ のいずれも予約済みなので、ここを $0$ にするのは不可能です。

したがって、今回の高橋くんは $5$ 点を得ることができます。

計算量

さて、方針は決まったもののそのまま実装すると計算量がマズそうです。例えば、木の頂点 $1$ から $N/2$ までが単純パスであり残りが頂点 $N/2$ の子であるとします。そして各頂点の得点がその頂点番号であるとすると、頂点 $N/2$ の子を順に $0$ にしていくことになりますが、ターンごとに遡る祖先の数が $1$ ずつ増えていくので、計算量が $\sum_{k=1}^{n/2} k$ くらいになってしまいます。

そこで、path compression という手法を用います。やることは単純で、一度の探索で通った各頂点の次の探索候補（本来は親）を最後に通った頂点のものに統一して書き換えるというものです。素集合データ構造（Union-Find）で見た人も多いと思います。実際、今回やっていることは行き先が同じ頂点の素集合（部分木の頂点集合）をその根を代表にして管理することだと言えます。

ちなみに、Nachiaさんのユーザー解説に出てくる Weighted Union Heuristic （所謂マージテク）も Union-Find に用いられている unoin by rank という手法と同様です。

Union-Find は path compression と unoin by rank を組み合わせることでメチャクチャ速くなっているらしいですが、どちらか片方だけでもN回探索の計算量が $O(N\log N)$ になるらしいです。

en.wikipedia.org

実はまだNachiaさんの解法を理解できていないのですが、なんとなく path compression の方が楽なのではないかという気がしています。

あとは $0$ にする頂点の順番を決めるのに適当なソートで $O(N\log N)$ 掛かるくらいなので、全体で $O(N \log N)$ です。

実装

$(A_i,i)$ のペアを優先度付きキューで並べ替えておき、DFSかBFSで各頂点の親を見つけておきます。

そして親とは別に次の探索候補を管理する配列を用意し、 $-1$ などの空きを表す値で初期化しておきます。なお、この配列の長さを $N+1$ として、（実際には存在しない）根の親を表す値を $N$ としておくと実装が楽になります。

あとはキューから順に取り出していき可能な限り探索をすれば良いです。好みですがキューに番兵として $(0,0)$ を置いています。

atcoder.jp

#include <bits/stdc++.h>
using namespace std;

int main() {
  int N;
  cin >> N;

  priority_queue<pair<int, int>> q;
  q.emplace(0, 0);
  for (int i = 1; i < N; i++) {
    int A;
    cin >> A;
    q.emplace(A, i);
  }

  vector<vector<int>> g(N);
  for (int i = 1; i < N; i++) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    g[u].emplace_back(v);
    g[v].emplace_back(u);
  }

  vector<int> parents(N);
  function<void(int, int)> dfs = [&](int x, int p) {
    parents[x] = p;
    for (int nx : g[x]) {
      if (nx != p) {
        dfs(nx, x);
      }
    }
  };
  // N: no parent
  dfs(0, N);

  // -1: non-zero, else: next candidate for non-zero
  vector<int> skip(N + 1, -1);
  // find the nearest non-zero ancestor
  function<int(int)> find_non_zero = [&](int x) {
    // path compression like DSU
    return skip[x] == -1 ? x : skip[x] = find_non_zero(skip[x]);
  };

  while (true) {
    auto [a, x] = q.top();
    q.pop();

    int non_zero = find_non_zero(parents[x]);
    if (non_zero == N) {
      cout << a << endl;
      break;
    }
    // replace with zero and set candidate
    skip[non_zero] = parents[non_zero];
  }
}