TakahiroNakamori

TakahiroNakamori

ナカモリタカヒロ

 
 

Subtree K-th Max

問題

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

記録

解法1

コード例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

コード例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;
}