Official

G - Game on Tree 2 Editorial by blackyuki


頂点 \(1\) を根とすると、駒は葉に向かって移動するだけなので、ゲーム終了時点での駒の位置によってそれまでに通った頂点とその中央値が一意に定まります。

全ての葉に対して、その頂点でゲームが終了したときの中央値を求めることさえできれば、木DPでこの問題が解けます。

よって DFS をしながら以下の処理が高速にできれば良いです。

  • 数の(多重)集合 \(S\)\(X\) を追加する。
  • \(S\) から \(X\) を削除する。
  • \(S\) の中央値を求める。

以下の2つの手法が考えられます。

1. BIT上の二分探索解

あらかじめ座標圧縮しておくことで、追加クエリと削除クエリは一点加算、中央値を求めるクエリは BIT 上の二分探索に置き換えられます。 計算量は \(O(NlogN)\) または \(O(Nlog^2N)\) です。

2. std::multiset解

小さい方半分と大きい方半分を管理する \(2\) つの multiset を用意します。常に要素数が半分ずつとなるようにすることで、クエリあたり \(O(logN)\) で処理できます。 以下の実装例はこちらの方針です。 計算量は \(O(NlogN)\) です。

実装例 (C++)

#include <bits/stdc++.h>
using namespace std;

multiset<int> S,T; // 常に S.size()==T.size() or S.size()==T.size()+1
void eval(){
    if(S.size()){
        auto itr = S.end(); itr--;
        T.insert(*itr); S.erase(itr);
    }
    while(S.size() < T.size()){
        S.insert(*T.begin()); T.erase(T.begin());
    }
}
int val(){
    auto itr = S.end(); itr--;
    if(S.size()==T.size()+1)return *itr;
    return (*itr+*T.begin())/2;
}
void erase(int x){
    auto itr = S.end(); itr--;
    if(*itr < x) T.erase(T.lower_bound(x));
    else S.erase(S.lower_bound(x));
}

int main(){
    int n; cin >> n;
    vector<int> A(n); for(int i = 0; i < n; i++) cin >> A[i];
    vector<int> dp(n);
    vector<vector<int>> g(n);
    for(int i = 0; i < n-1; i++){
        int a, b; cin >> a >> b; a--; b--;
        g[a].push_back(b); g[b].push_back(a);
    }
    function<void(int,int,int)> dfs=[&](int i,int p,int d){
        S.insert(A[i]); eval();
        int mi = 1000000005, ma = 0;
        for(int x : g[i]) if(x != p) {
            dfs(x, i, d+1);
            mi = min(mi,dp[x]); ma = max(ma,dp[x]);
        }
        if(ma == 0) dp[i] = val();
        else if(d&1) dp[i] = mi;
        else dp[i] = ma;
        erase(A[i]); eval();
    }; dfs(0,-1,0);
    cout << dp[0] << endl;
}

posted:
last update: