競技プログラミング日記

主に AtCoder の記事です

AtCoder Beginner Contest 239E

ABC239E

部分木に関する処理なので,自然に書けば DFS. 任意の \(i \in N\) に対して, \(K_{i}\) は \(K_{i} \leq 20\) と小さいので, 大きいほうから \(20\) 個すべて保持しておける.
各頂点に関して \(20\) 個まで保持しておいても, 常に空間計算量は \(O(20N)\) . 時間計算量の最悪ケースは \(O(20N log (20N))\), これは star 型の場合,中心に一番時間がかかる.

使っている記号,マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

vll x(1e5+10);
vvll to(1e5+10);
vvll ma(1e5+10, vll(20, -1)); // K-th max

void dfs(ll cu, ll pa){

  priority_queue<ll> pq;
  pq.push(x[cu]);

  fore(ne, to[cu]) if(ne != pa){
    dfs(ne, cu);
    fore(i, ma[ne]) pq.push(i);
  }

  rep(i,20) if(!pq.empty()){
    ma[cu][i] = pq.top();
    pq.pop();
  }

}

int main() {

  ll n,q; cin >> n >> q;
  rep(i,n) cin >> x[i];

  rep(i,n-1){
    ll a,b; cin >> a >> b;
    a--,b--;
    to[a].push_back(b);
    to[b].push_back(a);
  }

  dfs(0,-1);


  while(q--){
    ll v,k; cin >> v >> k;
    v--;
    k--;

    assert(ma[v].size() >= k);
    cout << ma[v][k] << endl;

  }

  return 0;
}