読者です 読者をやめる 読者になる 読者になる

AGC004 D - Teleporter

問題

agc004.contest.atcoder.jp

考察

K回の移動を行った際に全ての点が首都(1)にいないといけないという条件があります。
これにより、必ず1は1自身に戻るループにしないといけないことが分かります。
1以外の行き先だと、「全ての点が同着で1にいないといけない」 =「全ての点の向き先が同じでそれが1」となりますが、1が1以外に行くので矛盾してしまいます。

さて、上記考察により、全ての点がK回以前に1に辿り着けば題意を満たせます。
残念ながら、制約が厳しいため全ての点についてK回遡ってチェックしたりすることは難しそうです。

ここで、テレポートの移動をグラフ化してみます(1は自己ループに変更済み)。 「全ての町は首都に行ける」ため連結で、1を根とした木になることが分かります。
このグラフで、「どのノードからもK回以下の距離に1がある」ように辺を組み替えます。

そうすると、グラフを順に辿ってノードをK回程度移動した時に辺を繋ぎかえれば良さそうです。
なお、組み替えた辺の指す先は1以外である必要はありません。
ただし、以下の注意点があります。

  • 辺を差し替えるか判定する際は、もともと1に移動できるかに注意する
    • 1に移動できない場合、1への移動で1回消費するため、K-1個のノードごとに最後のノードを1へのパスにする
    • 1に移動できる場合は、K個のノードまで許容できる
  • ルートから貪欲に処理すると上手く行かないため、葉から判定する
    • 分岐点があると、最適な差し替えができません
    • 例えば、K=3で1 1 2 3 3 4 4は1が最適ですが、ルートから行うと2になります

上記条件を満たすため、dfsで葉っぱから要素数を判定してゆきます。
木の高さを判定するdfsと処理はほぼ同じですが、親(または自分)が1の時は高さをリセットする、という処理が必要です。
なお、以下のコードでは0-basedにしています。

コード

int N, K;
int ans = 0;
vvi r_edge;
 
int dfs(int v, int par) {
    int rank = 0;
    tr(it, r_edge[v]) {
        if (*it != par)
            chmax(rank, dfs(*it, v));
    }
 
    if (v == 0 || par == 0)
        rank = -1;
    else if (rank >= K - 1) {
        ans++;
        rank = -1;
    }
 
    return rank + 1;
}
 
 
int main() {
    cin >> N >> K;
 
    vi a(N);
    rep(i, N) {
        int t;
        cin >> t;
        t--;
        a[i] = t;
    }
 
    if (a[0] != 0)
        ans++;
    a[0] = 0;
 
    r_edge.resize(N);
    rep(i, N)
        r_edge[a[i]].push_back(i);
 
    dfs(0, 0);
 
    cout << ans << endl;
 
    return 0;
}