AtCoder Grand Contest 004 D - Teleporter

http://agc004.contest.atcoder.jp/tasks/agc004_d

問題概要

N個の町があり、1~Nまでの番号が振られている。それぞれの町には1台ずつテレポーターがあり、町iのテレポーターの転送先は町aiに設定されている。また初期状態では必ず何回かテレポーターをたどることで町1にたどり着けることが保証されている。

正整数Kが与えられたとき、「どの町を出発点としてもテレポーターをちょうどK回使うと町1にたどり着く」という状態を作りたい。最小で何個のテレポーターの転送先を変更するとこの条件を達成できるか求めよ。

2 <= N <= 105, 1 <= ai <= N, 1 <= K <= 109

解法

まず町1の転送先は必ず町1でなければならない。町1の転送先が町x (x!=1)でかつ町1→町1の移動がちょうどK回でできるとすると、町x→町1の移動は(K-1)回となり、結局町x→町xへの移動がちょうどK回となってしまうためである。

以上より町1の転送先を町1にすると、「初期状態では必ず何回かテレポーターをたどることで町1にたどり着けることが保証されている」という条件より、遷移を表したグラフは必ず木になることがわかる。この木において町1を根にした根付き木をGとすると、問題は以下のように言い換えられる。

Gのすべての頂点の深さをK以下にしたい。以下の操作を最小で何回行えば目的を達成できるか求めよ。【操作】頂点をひとつ選び、根が親になるよう辺を張り替える

ここで深さがKより大きいあるひとつの頂点に注目すると、自分・自分の1個上の頂点・自分の2個上の頂点・……・自分の(K-1)個上の頂点 のいずれかひとつに対して必ず1回は上の操作が行われている必要がある。また逆にある点を選んで操作を行ったとき、自分の(K-1)個下までの子供は条件を満たすことになる。ゆえにある特定の深さ>Kの点を深さK以下にしたいときには、(K-1)個上の親に操作を行うのが最適である。

ここまでを踏まえて、自分は以下のような貪欲をやった。まずマークされていない頂点の中で一番深いものを選び、その(K-1)個上の頂点の親を根に張り替える。その頂点の部分木に含まれる点をすべてマークする。これを深さ>Kの頂点がすべてマークされるまで続けたとき、親の張り替えを行った回数が答え。

感想

時間かかったけど自力で通せてよかった。ふと順位表を見たら当時の自分は1完だったらしく結構感慨深いものがありました。解法に関しては公式の解説だともっとスマートにやってる感じがあるんですが、葉の方から見てるという意味ではたぶん同じ、かなあ

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
void main() {
    auto s = readln.split.map!(to!int);
    auto N = s[0];
    auto K = s[1];
    auto A = readln.split.map!(to!int).array;
    auto edges = new int[][](N);
    foreach (i; 1..N) {
        auto a = A[i] - 1;
        auto b = i;
        edges[a] ~= b;
        edges[b] ~= a;
    }
    
    auto depth = new int[](N);
    void dfs(int n, int p, int d) {
        depth[n] = d;
        foreach (m; edges[n]) if (m != p) dfs(m, n, d + 1);
    }
    dfs(0, -1, 0);
 
 
    auto used = new bool[](N);
    void dfs2(int n, int lb, int ub) {
        used[n] = true;
        foreach (m; edges[n]) if (depth[m] >= lb && depth[m] <= ub && !used[m]) dfs2(m, lb, ub);
    }
 
    auto pq = new BinaryHeap!(Array!(Tuple!(int, int)), "a[1] < b[1]");
    foreach (i; 1..N) if (depth[i] > K) pq.insert(tuple(i, depth[i]));
    int ans = A[0] != 1;
 
    while (!pq.empty) {
        auto t = pq.front;
        auto n = t[0];
        pq.popFront;
        if (used[n]) continue;
        ans += 1;
        dfs2(n, min(depth[n], depth[n] - K + 1), depth[n]);
    }
 
    ans.writeln;
}