CF434E: Arkady and a Nobody-men

http://codeforces.com/contest/860/problem/E

HL分解やるだけで通ってしまった…
depthの浅い頂点から順番に見ていき、根まで+1をする更新を行うことで求められます。

segtreeでやったら遅かったのでbitにして1824ms/2000msでAC。

struct BIT { //0-origin!!! update and get just like lazy_segtree
	int n; vl bit0, bit1;
	void init(int mx) {
		n = mx;
		bit0 = vl(n + 1, 0); bit1 = vl(n + 1, 0);
	}
	BIT(int mx = 0) { init(mx); } 

	ll ga(vl& bit, int i) {
		ll s = 0;
		while(i > 0) { s += bit[i]; i -= i & -i; }
		return s;
	}
	void app(vl& bit, int i, ll x) {
		while(i <= n) { bit[i] += x; i += i & -i; }
	}
	void update(int a, int b, ll x) { //[a, b)
		a++;
		app(bit0, a, -x * (a - 1));
		app(bit1, a, x);
		app(bit0, b + 1, x * b);
		app(bit1, b + 1, -x);
	}
	ll get(int a, int b) { //[a, b)
		a++;
		return (ga(bit1, b) * b + ga(bit0, b)) 
			- (ga(bit1, a - 1) * (a - 1) + ga(bit0, a - 1));
	}
};

//////////////////////////////////////////////////////////////////////////HL BEGIN

struct HL {
	int n;
	vector<vi> to, cost, D;//D:group
	vi vcost, vd, vk, pv, dep;//vcost:edge cost,vd:D at,vk:num from parent,pv:group par,dep:depth
	//edge starts from (e:v->p) where v is top of D
	vector<BIT> S; //change 
	int root;

	HL(int mx = 0): n(mx), to(mx), cost(mx), vcost(mx), vd(mx), vk(mx), dep(mx) {}
	void add_edge(int a, int b, int c = 1) {
		to[a].pb(b); cost[a].pb(c);
		to[b].pb(a); cost[b].pb(c);
	}

	int dfs(int v, int depth = 0, int p = -1) {
		dep[v] = depth;
		vector<pi> rs;
		rep(i, 0, sz(to[v])) {
			int u = to[v][i];
			if(u == p) continue;
			vcost[u] = cost[v][i];
			rs.pb(pi(dfs(u, depth + 1, v), u));
		}
		if(sz(rs)) {
			sort(all(rs));
			vd[v] = vd[rs[0].sec]; pv[vd[v]] = p;
			D[vd[v]].pb(v); vk[v] = -sz(D[vd[v]]);
			return rs[0].fst - (sz(rs) != 1 && rs[0].fst == rs[1].fst);
		}
		else {
			vd[v] = sz(D); vk[v] = -1; D.pb(vi(1, v)); pv.pb(p);
			return 1;
		}
	}
	void init(int v = 0) {
		root = v;
		dfs(root);
		rep(i, 0, sz(D)) reverse(all(D[i]));
		rep(i, 0, sz(vk)) vk[i] += sz(D[vd[i]]);
		rep(i, 0, sz(D)) {
			S.pb(BIT(sz(D[i])));      
		}
	}
	int lca(int a, int b) {
		int p = vd[a]; vi ap(1, p);
		while(pv[p] != -1) ap.pb(p = vd[pv[p]]);
		reverse(all(ap)); ap.pb(-1);
		p = vd[b]; vi bp(1, p);
		while(pv[p] != -1) bp.pb(p = vd[pv[p]]);
		reverse(all(bp)); bp.pb(-2);
		int pi = 1; while(ap[pi] == bp[pi]) pi++;
		p = ap[pi - 1];
		int ai = vd[a] == p ? vk[a] : vk[pv[ap[pi]]];
		int bi = vd[b] == p ? vk[b] : vk[pv[bp[pi]]];
		return D[p][min(ai, bi)];
	}

	//////////////////////////////////////////////////////////////////////////////from here, change
	ll ga(int par, int v) {
		ll res = 0;
		while(true) {
			int nd = vd[v], nk = vk[v];
			if(nd == vd[par]) {
				return res + S[nd].get(vk[par], nk + 1);
			}
			else {
				res = res + S[nd].get(0, nk + 1);
			}
			v = pv[nd];
		}
	}
	void app(int par, int v, ll x) {
		while(true) {
			int nd = vd[v], nk = vk[v];
			if(nd == vd[par]) {
				S[nd].update(vk[par], nk + 1, x);
				return;
			}
			else {
				S[nd].update(0, nk + 1, x);
			}
			v = pv[nd];
		}
	}
	void query(int a) {
		app(root, a, 1);
		S[vd[a]].update(vk[a], vk[a] + 1, -1);
	}
	ll get(int a) {
		return ga(root, a);
	}
};

int N;
ll ans[MAX_N];
vector<int> D[MAX_N];
vector<int> G[MAX_N];

void solve() {
	cin >> N;
	HL hl(N);
	rep(i, 0, N) {
		int a;
		cin >> a; a--;
		if(a == -1) hl.root = i;
		else hl.add_edge(a, i);
	}
	hl.init(hl.root);
	rep(i, 0, N) {
		D[hl.dep[i]].pb(i);
	}
	int at = 1;
	while(!D[at].empty()) {
		rep(i, 0, sz(D[at])) {
			int v = D[at][i];
			hl.query(v);
		}
		rep(i, 0, sz(D[at])) {
			int v = D[at][i];
			ans[v] = hl.get(v);
		}
		at++;
	}
	rep(i, 0, N) cout << ans[i] << "\n";
}

HL分解ほんと定数倍軽いですね…