けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 004 D - Teleporter (800 点)

こういう系の問題、なかなか提出が怖くなるやつ。こういうのは背水の陣状態じゃないと躊躇してしまいそう...

問題へのリンク

問題概要

頂点数  N、辺数  N の有向グラフが与えられる。

  • どの頂点から出る辺数も 1
  • どの頂点から出発しても必ずノード 1 (root) にたどり着ける

という性質を持っている。このグラフの辺の行き先を最小回数入れ替えて、以下の条件を満たすようにせよ:

  • どの頂点から出発しても、ちょうど K 回の hop 回数を経ている場所がノード 1 (root) である

制約

  •  1 \le N \le 10^{5}

考えたこと

ほとんどツリーに近いグラフ!

根からぴゅっと飛び出てるけど、これはもし自己ループになっていなかったら、自己ループにしないとダメそう (自己ループになっていなかったらサイクルができるけど、サイクル上の点がすべて K-hop で根に行くことはできない)。

そうするともう、ツリーと変わらなくて、ツリー上の頂点を最小個数選んで、それらを根にくっつけることで、ツリーの高さを K 以下にする問題になる。

これはツリー DP でできそう。一見、 O(NK) かかってしまうように見えるのだが、dp 値を pair<int,int> 型にして、

dp[ v ] := 頂点 v を根とする部分木において、{選ぶ必要のある頂点数の最小値, またその場合での根から各選択頂点までの距離の最小値}

としてやることで  O(N) にすることができた。ここは色々やりようがありそう。

#include <iostream>
#include <vector>
using namespace std;

int N, K;
vector<vector<int> > G;

using pint = pair<int,int>;
vector<pint> dp;
pint rec(int v, int p) {
    for (auto nv : G[v]) {
        if (nv == p) continue;
        rec(nv, v);
    }
    pint res = {0, 1};
    for (auto nv : G[v]) {
        if (nv == p) continue;
        res.first += dp[nv].first;
        res.second = max(res.second, dp[nv].second + 1);
    }
    if (v != 0 && p != 0 && res.second >= K) ++res.first, res.second = 0;
    return dp[v] = res;
}

int main() {
    cin >> N >> K;
    int res = 0;
    G.clear(); G.resize(N);
    for (int i = 0; i < N; ++i) {
        int a; cin >> a; --a;
        if (i == 0 && a != 0) ++res;
        if (i != 0) {
            G[i].push_back(a);
            G[a].push_back(i);
        }
    }
    dp.assign(N, pint(N, N));
    rec(0, -1);
    cout << res + dp[0].first << endl;
}