AGC004 D - Teleporter

問題文

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

解法

首都が自己辺でなければならないのは，すぐ分かる．（ぱっとわからなくても実験すればわかるようになってる．）
あとは制約から，与えられたグラフが首都を根とする木になっていることがわかる．
首都が自己辺であれば，首都から距離 K 以内であればOKという問題に変わるので，その方向で考える．
この場合，葉から考えて，高さ K-1 の部分木を切っていくイメージでやるとよいことがわかる．
根から高さを数えてしまうと，根からの距離が K の節点に大量の葉がついているケースで死んでしまって，僕みたいに時間を溶かすのでだめ．
それがわかればDFSやるだけで求められる．

ソースコード

#include <bits/stdc++.h>
using namespace std;
 
constexpr int INF = 1e9;
 
using graph = vector<vector<int>>;
 
int N, K;
 
int dfs(int v, int prev, graph const& g, int& res) {
    int h = 0;
    for(auto to : g[v]) {
        h = max(h, dfs(to, v, g, res) + 1);
    }
    if(h == K - 1 && prev != 0 && v != 0) {
        res += 1;
        h = -1;
    }
    return h;
}
 
int main() {
    cin >> N >> K;
    graph g(N);
    int res = 0;
    for(int i = 0; i < N; ++i) {
        int a;
        cin >> a;
        a--;
        if(i == 0) {
            res += a != 0;
        } else {
            g[a].push_back(i);
        }
    }
    dfs(0, -1, g, res);
    cout << res << endl;
}

感想

根から高さ数えてて，死ぬケースもわかってたんだけど単純に葉から数えればいいという発想がなかった．
難しい問題解く前にこういう当たり前の（逆から考える）発想ができるようになりたい．