問題

E – Subtree K-th Max

プログラミング初級者から上級者まで楽しめる、競技プログラミングコンテストサイト「AtCoder」。オンラインで毎週開催プログラミングコンテストを開催しています。競技…

atcoder.jp

記録

• 2022/2/27　バーチャル参加。解法わからず。

解法1

• \(K \leq 20 \) なので各頂点は多い方から \(K \) 個だけ値を持てばいい。
• DFS で木の下から各頂点が持つ値（配列）を渡し、親の配列をソートして多い方から \(K \) 個だけに更新する。

コード例1

#include <bits/stdc++.h>
using namespace std;
#include <atcoder/all>
using namespace atcoder;
using ll = long long;
using P = pair<int, int>;
#define rep(i, n) for(int i = 0; i < (n); ++i)
// using mint = modint998244353;
// using mint = modint1000000007;
// const int mod = 1000000007;
// const ll INF = 1LL << 62;
// const int INF = 1001001001;

int N, Q;
int K = 20;
vector<vector<int>> edges;
vector<int> X;
vector<vector<int>> values;

vector<int> merge(vector<int> a, vector<int> b) {
  a.insert(a.end(), b.begin(), b.end());
  sort(a.rbegin(), a.rend());
  a.resize(K);
  return a;
}

void dfs(int current, int previous) {
  for(int next : edges[current]) {
    if(next != previous) {
      dfs(next, current);
      values[current] = merge(values[current], values[next]);
    }
  }
}

int main() {
  cin >> N >> Q;
  X.resize(N);
  edges.resize(N);
  values.resize(N);

  rep(i, N) {
    cin >> X[i];
    values[i].push_back(X[i]);
  }

  for(int i = 0; i < N - 1; ++i) {
    int A, B;
    cin >> A >> B;
    A--;
    B--;
    edges[A].push_back(B);
    edges[B].push_back(A);
  }

  dfs(0, -1);

  while(0 < Q) {
    int v, k;
    cin >> v >> k;
    v--;
    k--;
    cout << values[v][k] << endl;
    Q--;
  }

  return 0;
}

解法2

• \(K \leq 20 \) なので各頂点は多い方から \(K \) 個だけ値を持てばいい。
• オイラーツアー 辿った頂点の順番を把握する。
• 頂点 \(V\) に対する部分木に対応する区間 \([L,R)\) を把握する。
• セグメントツリーを \(K\) 回行う（最大値を求めて、\(-1\) に更新して除外する…を繰り返して終わったら戻す）。
• 戻せるようにするために、セグツリーはpairで行うなどの工夫が必要。

コード例2

#include <bits/stdc++.h>
using namespace std;
#include <atcoder/all>
using namespace atcoder;
using ll = long long;
using P = pair<int, int>;
#define rep(i, n) for(int i = 0; i < (n); ++i)
// using mint = modint998244353;
// using mint = modint1000000007;
// const int mod = 1000000007;
// const ll INF = 1LL << 62;
const int INF = 1001001001;

int N, Q;
vector<vector<int>> edges;
vector<int> X;
vector<int> depth;
// 通った点
vector<int> tour;
// 通った辺
vector<int> tourEdge;
vector<int> discovery;
vector<int> finishing;

// オイラーツアー
void f(int a, int d) {
  depth[a] = d;
  tour.push_back(a);
  tourEdge.push_back(a);
  int n = edges[a].size();
  for(int i = 0; i < n; ++i) {
    int next = edges[a][i];
    if(depth[next] == -1) {
      f(next, d + 1);
      tour.push_back(a);
    }
  }
  tourEdge.push_back(-a);
}

// SegmentTree
P op(P a, P b) {
  if(a.first < b.first) {
    return b;
  } else {
    return a;
  }
}

P e() {
  return make_pair(-1, -1);
}

int target;
bool f(int v) {
  return v < target;
}

int main() {
  cin >> N >> Q;
  X.resize(N);
  edges.resize(N);
  discovery.resize(N, INF);
  finishing.resize(N, -1);
  depth.resize(N, -1);

  rep(i, N) { cin >> X[i]; }

  for(int i = 0; i < N - 1; ++i) {
    int A, B;
    cin >> A >> B;
    A--;
    B--;
    edges[A].push_back(B);
    edges[B].push_back(A);
  }

  f(0, 0);

  segtree<P, op, e> seg(N * 2 + 2);

  int n = tour.size();
  rep(i, n - 1) {
    if(0 <= tourEdge[i]) {
      seg.set(i, make_pair(X[tour[i]], i));
      X[tour[i]] = -1;
    } else {
      seg.set(i, make_pair(-1, -1));
    }
  }

  rep(i, n) {
    int index = tour[i];
    discovery[index] = min(i, discovery[index]);
    finishing[index] = max(i, finishing[index]);
  }

  rep(i, Q) {
    int V, K;
    cin >> V >> K;
    int sl = discovery[V - 1];
    int sr = finishing[V - 1] + 1;
    vector<P> data;

    rep(j, K - 1) {
      P val = seg.prod(sl, sr);
      int index = val.second;
      data.push_back(val);
      seg.set(index, make_pair(-1, index));
    }

    P ans = seg.prod(sl, sr);
    cout << ans.first << endl;
    if(1 < K) {
      for(auto j : data) {
        seg.set(j.second, make_pair(j.first, j.second));
      }
    }
  }

  return 0;
}