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

問題

code-festival-2015-final-open.contest.atcoder.jp

解法

公式解説がとても分かりやすいので解説はほとんど省略する。
CODE FESTIVAL 2015 解説

頂点vを最大の高さにして残すと決め打ちしたとき何本の糸を切ればいいのか求めたい。これは、根と頂点vを結ぶパスから生えてる部分木のうち、高さheight[v]より大きな風船を持つようなものの数を数えればいい。

頂点のラベルには風船の高さ、部分木のラベルには風船の最大の高さが書かれている。

この例の場合height[v]=3なので4つ部分木を切り離せばいい。

何本切ればいいのかはBIT（あるいはSegmentTree）をごにょごにょしながらDFSすれば計算できる。

#include <bits/stdc++.h>
#define GET_MACRO(a, b, c, NAME, ...) NAME
#define rep(...) GET_MACRO(__VA_ARGS__, rep3, rep2)(__VA_ARGS__)
#define rep2(i, a) rep3 (i, 0, a)
#define rep3(i, a, b) for (int i = (a); i < (b); i++)
#define repr(...) GET_MACRO(__VA_ARGS__, repr3, repr2)(__VA_ARGS__)
#define repr2(i, a) repr3 (i, 0, a)
#define repr3(i, a, b) for (int i = (b) - 1; i >= (a); i--)
#define chmin(a, b) ((b) < a && (a = (b), true))
#define chmax(a, b) (a < (b) && (a = (b), true))
using namespace std;
typedef long long ll;
 
const int inf = 1e9;
 
struct BIT {
    int size;
    vector<ll> bit;
    BIT(int size) : size(size), bit(size + 1) {}
    void update(int k, ll x) {
        while (k <= size) {
            bit[k] += x;
            k += k & -k;
        }
    }
    ll query(int k) {
        ll res = 0;
        while (k > 0) {
            res += bit[k];
            k -= k & -k;
        }
        return res;
    }
};
 
struct edge { int v, w; };
vector<edge> g[101010];
int max_height[101010];
int height[101010];
BIT bit(303030);
int ans[303030];
 
void dfs(int curr, int prev = -1) {
    max_height[curr] = height[curr];
    for (auto e : g[curr]) if (e.v != prev) {
        height[e.v] = height[curr] + e.w;
        dfs(e.v, curr);
        chmax(max_height[curr], max_height[e.v]);
    }
}
 
void dfs2(int curr, int prev = -1) {
    for (auto e : g[curr]) if (e.v != prev) {
        bit.update(max_height[e.v], 1);
    }
    int cand = bit.query(303030) - bit.query(height[curr]);
    chmin(ans[height[curr]], cand);
    for (auto e : g[curr]) if (e.v != prev) {
        bit.update(max_height[e.v], -1);
        dfs2(e.v, curr);
        bit.update(max_height[e.v], 1);
    }
    for (auto e : g[curr]) if (e.v != prev) {
        bit.update(max_height[e.v], -1);
    }
}
 
int main() {
    int n;
    cin >> n;
    vector<int> l(n), p(n);
    rep (i, n) scanf("%d", &l[i]);
    rep (i, 1, n) scanf("%d", &p[i]);
    rep (i, 1, n) {
        g[i + 1].push_back({p[i], l[i]});
        g[p[i]].push_back({i + 1, l[i]});
    }
    n++;
    g[0].push_back({1, l[0]});
    g[1].push_back({0, l[0]});
    dfs(0);
    int q;
    cin >> q;
    vector<int> h(q);
    rep (i, q) scanf("%d", &h[i]);
 
    vector<int> ch;
    rep (i, n) ch.push_back(max_height[i]);
    rep (i, q) ch.push_back(h[i]);
    rep (i, q) ch.push_back(height[i]);
    sort(ch.begin(), ch.end());
    rep (i, n) {
        max_height[i] = lower_bound(ch.begin(), ch.end(), max_height[i]) - ch.begin();
        height[i] = lower_bound(ch.begin(), ch.end(), height[i]) - ch.begin();
    }
    rep (i, q) {
        h[i] = lower_bound(ch.begin(), ch.end(), h[i]) - ch.begin();
    }
 
 
    rep (i, 303030) ans[i] = inf;
    dfs2(0);
    rep (i, 303030) if (ans[i] == inf) ans[i] = -1;
    rep (i, q) printf("%d\n", ans[h[i]]);
    return 0;
}

コメント

本番中に解きたい問題だった。