mayoko’s diary

プロコンとかいろいろ。

CODE FESTIVAL 2015 決勝 I - 風船ツリー

解法

やっぱり解説がわかりやすいと思いますね…

前準備として, depth[i] = (i 番目の頂点の高さ), maxDepth[i] = (i を根とする部分木に属する頂点の高さの中で最大のもの) という値を簡単な木 DP で求めます。

で, 問題の目的を達成するためには, 「頂点 v を残す(頂点 v の高さを最大とするように紐をほどく)ために最低限ほどかないといけない紐の数」を求めれば, 各クエリ H[q] が来た時, 高さが H[q] となる頂点のなかで紐をほどく必要のある頂点が最小のものを求めれば良いだけになります。

ということで, 「頂点 v を残す(頂点 v の高さを最大とするように紐をほどく)ために最低限ほどかないといけない紐の数」を求めます。これをやるために, (少なくとも自分にとっては)かなり頭の良い BIT + dfs を使います。

頂点 v で dfs をした時,
v の各子 ch について, BIT に maxDepth[ch] を追加
v を残す際にほどかないといけない紐の数を求める。これは, BIT.sum(depth[v], N) で求められる。
v から ch に dfs する場合, BIT から maxDepth[ch] を引く
戻ってきたら再び BIT に maxDepth[ch] を足す

このようにすると, 頂点 v にたどり着いた時点では, v と別の分岐をしている頂点については maxDepth を調べれば良いことを表現していて v の子になっているものについて刈り取る頂点の数もわかるのでハッピーです。よく出来過ぎでしょこれ。

// 0-based Binary Indexed Tree
// 数え上げ用
constexpr int mod = 1e9+7;
template<typename T> struct BIT {
    int max;
    vector<T> bit;
    BIT(int max) : max(max) {bit.resize(max+1);}
    // [0, i)
    T sum(int i) {
        T s = 0;
        while (i > 0) {
            (s += bit[i]) %= mod;
            i ^= i&-i;
        }
        return s;
    }
    // 0-basedな座標iに値xを追加する
    void add(int i, T x) {
        ++i;
        while (i <= max) {
            (bit[i] += x) %= mod;
            i += i&-i;
        }
    }
    // [a, b)
    T sum(int a, int b) {
        return sum(b)-sum(a);
    }
    // sum(0, i) >= wとなる最小のiを求める 存在しなければmaxを返す
    int lb(T w) {
        if (w <= 0) return 0;
        int k = 1;
        while (k <= max) k <<= 1;
        int i = 0;
        for (; k > 0; k >>= 1) if (i+k <= max && bit[i+k] < w) {
            w -= bit[i+k];
            i += k;
        }
        return i+1;
    }
};


const int MAXN = 100010;
ll L[MAXN];
ll depth[MAXN];
int maxDepth[MAXN], need[MAXN], ans[MAXN];
vector<int> G[MAXN];
BIT<int> bit(MAXN);

void getDepth(int v, int p, ll d) {
    depth[v] = d;
    for (int ch : G[v]) {
        if (ch == p) continue;
        getDepth(ch, v, d+L[ch]);
    }
}

int getMaxDepth(int v, int p) {
    maxDepth[v] = depth[v];
    for (int ch : G[v]) {
        if (ch == p) continue;
        maxDepth[v] = max(maxDepth[v], getMaxDepth(ch, v));
    }
    return maxDepth[v];
}

void dfs(int v, int p) {
    for (int ch : G[v]) {
        if (ch == p) continue;
        bit.add(maxDepth[ch], 1);
    }
    need[v] = bit.sum(depth[v]+1, MAXN-1);
    for (int ch : G[v]) {
        if (ch == p) continue;
        bit.add(maxDepth[ch], -1);
        dfs(ch, v);
        bit.add(maxDepth[ch], 1);
    }
    for (int ch : G[v]) {
        if (ch == p) continue;
        bit.add(maxDepth[ch], -1);
    }
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N;
    cin >> N;
    for (int i = 0; i < N; i++) cin >> L[i];
    for (int i = 1; i < N; i++) {
        int p;
        cin >> p;
        p--;
        G[i].push_back(p);
        G[p].push_back(i);
    }
    // 実際の深さを求める
    getDepth(0, -1, L[0]);
    // 変換する
    map<ll, int> M;
    for (int i = 0; i < N; i++) M[depth[i]] = 0;
    {
        int k = 0;
        for (auto& p : M) p.second = ++k;
    }
    for (int i = 0; i < N; i++) depth[i] = M[depth[i]];
    // maxDepth を求める
    getMaxDepth(0, -1);
    // solve
    dfs(0, -1);
    for (int i = 0; i < MAXN; i++) ans[i] = MAXN;
    for (int i = 0; i < N; i++) {
        ans[depth[i]] = min(ans[depth[i]], need[i]);
    }
    int Q;
    cin >> Q;
    while (Q--) {
        ll h;
        cin >> h;
        if (M[h] == 0) {
            cout << -1 << endl;
        } else {
            cout << ans[M[h]] << endl;
        }
    }
    return 0;
}