Official

G - Game on Tree 2 Editorial by en_translator


Let us assume Vertex \(1\) as the root. Then the piece can only move towards leaves, so the final position of the piece at the end of the game uniquely determines the set of vertices passed through so far and the median.

Once we find for every leaf the median after finishing the game at the vertex, we can solve the problem with a DP (Dynamic Programming) on the tree.

Therefore, during DFS, we want to process the following operations fast:

  • Add \(X\) to a (multi-)set \(S\) of integers.
  • Remove \(X\) from \(S\).
  • Find the median of \(S\).

The following two approaches are possible.

1. Binary search on BIT (Binary Indexed Tree)

By coordinate-compressing them beforehand, we can process the insertion query and removal query as an addition to one element, and the median query as a binary search on BIT. The time complexity is \(O(NlogN)\) or \(O(Nlog^2N)\).

2. Solution using std::multiset

Prepare two multisets managing the smaller half and the larger half. By maintaining the number of elements of each multiset to be half the total, we can process in an \(O(logN)\) time each. The following sample code adopts this approach. The time complexity is \(O(NlogN)\).

Sample code (C++)

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

multiset<int> S,T; // Always 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: