Codeforces Round #359 (Div. 2) - D. Kay and Snowflake

Problem Statement:
http://codeforces.com/contest/686/problem/D

Summary:
Given a tree of size n <= 300000, answer q <= 300000 queries of the form: find the centroid of subtree v. Centroid of subtree T is defined as a node in T such that when removed the resulting set of disjoint trees have sizes at most |T|/2 (half the size of initial subtree).



Solution:
Focus on a tree T and try to locate the centroid. The centroid of T must be the root of a subtree S that with size |S| >= |T|/2 (because otherwise, if S has size < |T|/2, then the remaining connected component must have size |T| - |S| >= |T|/2). Starting from root of T, by following the path of the largest child down the tree, we will reach S (on each node on this path, there will exactly be one child with size >= |T|/2). Furthermore, by definition, children of S must have sizes less than |T|/2, hence S is the smallest subtree in this path with size >= |T|/2.

Based on above observation, we can come up with a linear solution in n and q. First we compute the centroid for every subtree, and then we answer the queries. We will employ a simple recursive solution to compute a centroid of subtree rooted at u, call it T(u). For each child of u, compute its centroid. Then take note of the largest child v. We know the centroid of v, say w. Let's iterate through the path from w to u, starting from w, and check if current node has size >= |T(u)|/2. If so, we have found the centroid of u! If not, we one step up the path (i.e. go to the parent of current node). Note that it is possible that u is the only node that has size >= |T(u)|/2, which means that u is the centroid of its own subtree. If we proceed this computation in bottom up order, the whole recursion will take O(n) time (This is because the centroid of a tree cannot be lower than the centroid of its subtrees, hence when we consider the path w-v-u, the nodes that we have visited will not be visited more than once up the recursion tree).

Implementation:

#include <cstdio>
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;
vector<vector<int>> adj;
int n, q;
int par[300010];
int sz[300010], centroid[300010];
void ComputeSize(int u) {
  for (int i=0;i<adj[u].size();++i){
    ComputeSize(adj[u][i]);
    sz[u] += sz[adj[u][i]];
  }
  sz[u]++;
}
void ComputeCentroid(int u) {
  if (sz[u]==1) {
    centroid[u] = u;
    return;
  }
  int largest = adj[u][0];
  for (int i=0;i<adj[u].size();++i){
    int v = adj[u][i];
    ComputeCentroid(v);
    largest = sz[largest] < sz[v] ? v : largest;
  }
  int half = sz[u] - sz[u]/2;
  int cur = centroid[largest];
  while (cur != u) {
    if (sz[cur] < half) {
      cur = par[cur];
    } else break;
  }
  centroid[u] = cur;
}
int main() {
  scanf("%d%d",&n,&q);
  adj.resize(n+2);
  for(int i=1;i<n;++i){
    scanf("%d",&par[i]);
    par[i]--;
    adj[par[i]].push_back(i);
  }
  par[0] = -1;
  ComputeSize(0);
  ComputeCentroid(0);
  int u;
  for (int i=0;i<q;++i){
    scanf("%d",&u);
    printf("%d\n",centroid[u-1]+1);
  }
  return 0;
}