CodeFestival2015本戦 I問題 風船ツリー

んーなんかF,G,Hよりは本番中にできる可能性のある問題だったなぁ。
今回も解説スライドがわかりやすいのでみちくり。

• 座標が1e9までだが、高さの上下関係しか興味がないから1~Nに落とせる。
• 最終的にvの高さを超えるものは根本で切れ！
• つまり求めるのは0->vに行くまでの直接の子孫でvの高さを超えるものの数がほしいことに気づけ。
• それはBinaryIndexTreeをうまく使えば簡単に求まることに気づけ。

という感じ。
実装はそこまで難しくないし、F以降だと一番直感的で手をつければ良かったなぁと思った。

struct BinaryIndexedTree{
    int N;
    std::vector<int> bit; //1-indexed つまり1~Nを使う
    BinaryIndexedTree(int n): N(n), bit(n+1, 0){}
    void add(int a, int w){
        for(int x = a; x <=N; x+=x&(-x)) bit[x] += w;
    }
    int sum(int a){ //1~aまでsum
        int ret = 0;
        for(int x = a; x>0; x-=x&(-x)) ret += bit[x];
        return ret;
    }
};

const int MAXN = 100100;

int N;
int L[MAXN], P[MAXN];
int maxH[MAXN];
int H[MAXN];
int cut[MAXN];
vector<int> G[MAXN];
BinaryIndexedTree BIT(MAXN);
map<int, int> mp;
int ans[MAXN];

//風船それぞれの高さを求める
void dfs1(int v, int now)
{
    H[v] = now;
    for(auto u : G[v]){
        dfs1(u, now+L[u]);
    }
}

//風船それぞれから最大で行ける高さを求める
int dfs2(int v)
{
    int mx = H[v];
    for(auto u : G[v]){
        mx = max(mx, dfs2(u));
    }
    return maxH[v] = mx;
}

void dfs3(int v)
{
    for(auto u : G[v]){
        BIT.add(maxH[u], 1);
    }
    cut[v] = BIT.sum(MAXN) - BIT.sum(H[v]);
    for(auto u : G[v]){
        BIT.add(maxH[u], -1);
        dfs3(u);
        BIT.add(maxH[u], 1);
    }
    for(auto u : G[v]){
        BIT.add(maxH[u], -1);
    }
}


int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    cin >> N;
    rep(i, N) cin >> L[i];
    rep(i, N-1) cin >> P[i+1], P[i+1]--;

    for(int i=1; i<N; i++){
        G[P[i]].push_back(i);
    }

    dfs1(0, L[0]);

    //高さを圧縮
    rep(i, N) mp[H[i]] = 0;
    {int h=1; for(auto& m : mp) m.second = h, h++;}
    rep(i, N) H[i] = mp[H[i]];

    dfs2(0);

    dfs3(0);

    rep(i, MAXN) ans[i] = 1e9;
    rep(i, N){
        ans[H[i]] = min(ans[H[i]], cut[i]);
    }

    int Q;
    cin >> Q;
    rep(i, Q){
        int h;
        cin >> h;
        if(mp[h] == 0){
            cout << -1 << endl;
        }else{
            cout << ans[mp[h]] << endl;
        }
    }
}