CodeChef February Challenge 2019: Yet Another Tree Problem - 日々ｄｒｄｒする人のメモ

CodeChef February Challenge 2019 の問題: Yet Another Tree Problem (Code: TRDST)
問題ページ: https://www.codechef.com/FEB19A/problems/TRDST

問題概要

頂点に1から ${N}$ までの番号が付いた ${N}$ 頂点の木が与えられる。
また、頂点 ${u}$ と頂点 ${v}$ の距離を ${d(u, v)}$ とする。

${K_1, K_2, ..., K_N}$ が与えられるので、各頂点 ${i}$ について、 ${d(i, v) > D_i}$ となる頂点 ${v}$ が ${K_i}$ 個以上になるような最大の ${D_i}$ を求めよ。

制約

${1 \le N \le 10^5}$
各iについて ${1 \le K_i \le N - 1}$

部分点制約

${1 \le N \le 10^3}$

解法

部分点解法

各頂点iから探索して、 ${d(i, v) \le D_i}$ となる頂点の個数が ${N - K_i}$ 個以下になる最大の ${D_i}$ を解とする。

計算量は ${O(N^2)}$

満点解法

木の重心分解の上でクエリ処理ができることを知らずに、結構悩んでた。

この問題は、二分探索と、木の重心分解上のクエリ処理で解ける。

まず、ある頂点iについて、 ${D_i}$ を固定してその時の ${d(i, v) \le D_i}$ となる頂点個数を計算することを考える。
これは木の重心分解をして、クエリ処理をすることで一回あたり ${O(\log N)}$ で実現できる。

1. 前計算

前計算として、各部分木において重心となる頂点 ${v}$ について以下の値を計算しておく。

${A(v, d)}$ : 頂点 ${v}$ から距離 ${d}$ 以下となる頂点 ${w}$ の個数
${B(v, x, d)}$ : 頂点 ${v}$ に隣接する頂点 ${x}$ について、頂点 ${x}$ を通って到達できる頂点 ${w}$ のうち、距離 ${d}$ 以下となる頂点 ${w}$ の個数

この前計算の計算量と空間量はそれぞれ ${O(N \log N)}$ になる。

2. クエリ処理

クエリ処理において、各頂点 ${i}$ について ${D_i}$ 以下の距離にある頂点を数えることを考える。

この時、頂点 ${i}$ を含む各部分木の重心頂点について、以下の計算を行う。

頂点iが重心となる部分木において、
- 頂点iから距離 ${D_i}$ 以下にある頂点個数 ${A(i, D_i)}$ を解に加える
頂点v ()が重心となる部分木において、
- 頂点vから距離 ${D_i - d(i, v)}$ 以下にある頂点の個数 ${A(v, D_i - d(i, v))}$ を解に加える
- このままでは小さい部分木と重複して数えているため、頂点vから頂点iに到達するために通る、頂点vの隣接頂点を ${x}$ とした時、 ${B(v, x, D_i - d(i, v))}$ を解から引く

この計算によって、距離 ${D_i}$ 以下の頂点の個数を計算量 ${O(\log N)}$ で計算できる。

あとは二分探索で ${d(i, v) \le D_i}$ となる頂点の個数が ${N - K_i}$ 個以下になる最大の ${D_i}$ を求めればよい。

計算量は ${O(N \log^2 N)}$

実装

提出コード(部分点解法, PyPy3): Solution: 22861776 | CodeChef

木の重心分解を書きなれてないので雑な実装になってる...。

提出コード(C++14): Solution: 23017576 | CodeChef

#define N 100003
vector<int> g[N];
int n, k[N];

int root;
vector<int> g0[N];
int parent[N], level[N];

// ===== 木の重心分解
bool c_used[N];
int sz[N];
int count_dfs(int v, int p) {
  int c = 1;
  repe(w, g[v]) {
    if(w == p || c_used[w]) continue;
    c += count_dfs(w, v);
  }
  return sz[v] = c;
}

int search_centroid(int v, int p, int s) {
  if((s - sz[v])*2 > s) {
    return -1;
  }
  bool ok = true;
  int x;
  repe(w, g[v]) {
    if(w == p || c_used[w]) continue;

    int x = search_centroid(w, v, s);
    if(x != -1) {
      return x;
    }
    if(sz[w]*2 > s) {
      ok = false;
    }
  }
  return ok ? v : -1;
}

// mp[u][v] = i:
//  重心頂点uから、その重心で分割される部分木の重心頂点vに
//  到達するために通る隣接頂点は頂点g[u][i]
unordered_map<int, int> mp[N];

void centroid() {
  queue<int> que;

  int s = count_dfs(0, -1);
  int x = search_centroid(0, -1, s);
  que.push(x);
  c_used[x] = true;
  root = x;
  parent[x] = -1;
  level[x] = 0;

  while(!que.empty()) {
    int v = que.front(); que.pop();
    rep(i, g[v].size()) {
      int w = g[v][i];
      if(c_used[w]) {
        g0[v].push_back(-1);
        continue;
      }

      int s = count_dfs(w, -1);
      int x = search_centroid(w, -1, s);
      g0[v].push_back(x);
      parent[x] = v;
      level[x] = level[v] + 1;
      mp[v][x] = i;

      que.push(x);
      c_used[x] = true;
    }
  }
}
// ==== 木の重心分解ここまで

// ds[u][d] = cnt:
//  頂点uが重心となる部分木において、
//  頂点uから距離d以下に存在する頂点の個数cnt
vector<int> ds[N];

// de[u][i][d] = cnt:
//  頂点uが重心となる部分木において、
//  隣接する頂点g[u][i]を通って到達できる頂点の中で、
//  頂点uから距離d以下に存在する頂点の個数cnt
vector<vector<int> > de[N];

// vd[u][v] = d: 頂点uが重心となる部分木において、頂点uから頂点vまでの距離
unordered_map<int, int> vd[N];

int dd[N]; // これはbfs用の一時変数
// 前計算処理
void prepare() {
  rep(i, n) {
    // 頂点iが重心となる部分木において、頂点iから頂点vまでの距離を計算
    // - vd[i] と ds[i] の内容を計算
    int l = level[i];
    queue<P> que;
    dd[i] = 0;
    vd[i][i] = 0;
    que.push(P(i, -1));
    // cur: 今見ている距離, c: 現在の頂点個数
    int cur = 0, c = 0;
    while(!que.empty()) {
      P &p = que.front();
      int v = p.first, prv = p.second;
      que.pop();
      if(cur < dd[v]) {
        // cur < dd[v] となった時、距離curの位置にある頂点はもう出現しないので、
        // 距離d以内に存在する頂点の個数はc
        ds[i].push_back(c);
        ++cur;
      }
      ++c;
      repe(w, g[v]) {
        if(w == prv || level[w] <= l) continue;
        vd[i][w] = dd[w] = dd[v] + 1;
        que.push(P(w, v));
      }
    }
    ds[i].push_back(c);

    // 頂点iが重心となる部分木において、
    // 頂点iに隣接する頂点xに対し de[i][j] (x = g[i][j]) を求める
    queue<int> que0;
    repe(x, g[i]) {
      vector<int> d0;
      if(level[x] <= l) {
        de[i].push_back(d0);
        continue;
      }
      // cur: 今見ている距離, c: 現在の頂点個数 (頂点iは含めない)
      int cur = 1, c = 0;
      que0.push(x);
      d0.push_back(0);
      while(!que0.empty()) {
        int v = que0.front(); que0.pop();
        if(cur < dd[v]) {
          d0.push_back(c);
          ++cur;
        }
        ++c;
        repe(w, g[v]) {
          if(level[w] <= l || dd[w] <= dd[v]) continue;
          que0.push(w);
        }
      }
      d0.push_back(c);
      de[i].push_back(d0);
    }
  }
}

// 頂点vから距離r以内にある頂点の個数を求める
int node_count(int v, int r) {
  int res = 0, prv = -1;
  int w = v;
  // 頂点iを含む部分木を一つずつ上がっていく
  while(w != -1) {
    // d: 重心となる頂点wから頂点vの距離
    int d = vd[w][v];
    if(d <= r) {
      // ここでは、頂点wから距離(r-d)以内にある頂点個数を求めて結果に加える
      int d1 = r - d;
      res += ds[w][min(d1, (int)ds[w].size()-1)];

      if(prv != -1) {
        // 頂点wから頂点vが存在する方向の隣接頂点を通って到達できる頂点のうち、
        // 距離(r-d)以内にある頂点個数を求めて結果から引く
        // (少し見にくかったので実際のコードから少し修正)
        int i = mp[w][prv];
        res -= de[w][i][min(d1, (int)de[w][i].size()-1)];;
      }
    }
    prv = w; w = parent[w];
  }
  return res;
}

void solve() {
  prepare();

  // クエリ処理
  rep(i, n) {
    // 距離D_iに関する二分探索
    int base = n-k[i];
    int left = 0, right = n;
    while(left + 1 < right) {
      int mid = (left + right) >> 1;
      int c = node_count(i, mid);
      if(c <= base) {
        left = mid;
      } else {
        right = mid;
      }
    }
    printf((i ? " %d" : "%d"), left);
  }
  printf("\n");
}

int main() {
  scanf("%d", &n);
  rep(i, n) scanf("%d", k+i);

  rep(i, n-1) {
    int u, v; scanf("%d %d", &u, &v); --u; --v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  centroid();

  solve();
  return 0;
}