Heaps
Contest: Codeforces Round 471 (Div. 2)
Difficulty rating: 2600
Tags: dp, trees

You're given a tree with n vertices rooted at 1.

We say that there's a k-ary heap of depth m located at u if the following holds:

- For m = 1 u itself is a k-ary heap of depth 1.
- For m > 1 vertex u is a k-ary heap of depth m if at least k of its children are k-ary heaps of depth at least m - 1.

Denote dpk(u) as maximum depth of k-ary heap in the subtree of u (including u). Your goal is to compute $$\sum_{k=1}^{n}\sum_{u=1}^{n}dp_k(u)$$.

Input format:
The first line contains an integer n denoting the size of the tree (2 ≤ n ≤ 3·105).

The next n - 1 lines contain two integers u, v each, describing vertices connected by i-th edge.

It's guaranteed that the given configuration forms a tree.

Output format:
Output the answer to the task.

Constraints:
Time limit: 2.5 s
Memory limit: 512.0 MB

Note:
Consider sample case one.

For k ≥ 3 all dpk will be equal to 1.

For k = 2 dpk is 2 if $$u \in 1,3$$ and 1 otherwise.

For k = 1 dpk values are (3, 1, 2, 1) respectively.

To sum up, 4·1 + 4·1 + 2·2 + 2·1 + 3 + 1 + 2 + 1 = 21.

Example input:
4
1 3
2 3
4 3

Example output:
21

Example input:
4
1 2
2 3
3 4

Example output:
22

Editorial:
Denote heapk(u) as maximum depth of k-ary heap rooted at u, and dpk(u) as maximum depth of k-ary heap in the subtree of u (including u). Let vi — children of vertex u; sort them in order of descending heapk{vi}. Then heapk(u) = heapk(vk) + 1, and dpk(u) = max(heapk(u), maxj(dpk(vj))). So we can calculate dp for fixed k in $$O(n \log n)$$ or in O(n), if we will use nth_element, so we can solve the problem in O(n2).
For simplicity, denote dpk(1) as pk. Let's suppose that k > 1. Note three facts:
• pk ≤ logk(n);
• dpk(u) ≤ 2 when $$k \geq \sqrt{n}$$.
• heapk(u) ≥ heapk + 1(u).
So we can solve the problem in $$O(n\sqrt{n})$$. Let's solve task in O(n) when $$k \leq \sqrt{n}$$. For other k's let's go in descending order, and set the value of dp to 2 and push this value up to the root for each vertices, that have exactly k children.
Note, that the total number of vertices visited in pushes does not exceed n (Because if we arrive at vertex with dp = 2, it is useless to go up, because everything there has already been updated. Each vertex will be added exactly one time, so complexity of this part is O(n).
Let's use this idea to solve the problem in $$n^{\frac{1}{3}}$$. For each $$k \leq n^{\frac{1}{3}}$$ let's solve in O(n), and for $$k \geq \sqrt{n}$$ let's use the idea above. And when $$n^{\frac{1}{3}} \leq k \leq \sqrt{n}$$, dpk(u) does not exceed 3. Let f3(u) will be minimal k, such that heapk(u) is equal to 3. By definition, u will have at least k children, which are the heaps of depth 2, that is also vertices with at least k children. Let's sort v by number of their children; Then answer for u will be maximal k, such that children(vk) ≥ k. Let's precalculate it, and when let's go in descending order by k pushing up value 3 from each u with f3(u) = k. The total number of vertices visited in pushes does not exceed 2·n (At worst case, in each vertex will be pushed at first with value 2, and then with value 3). So we can solve the problem in $${\mathcal O}(n^{{\frac {4}{3}}})$$, which is better but still not enough.
Let val(u, depth) — maximal k > 1, such that vertex u contain k-ary heap of depth depth (or  - 1, if there are no such vertex). This dp have $$n \log n$$ states; To recalculate val(u, depth) we need to sort val(vi, depth - 1) by descending order, and find maximal k, such that val(vk, depth - 1) ≥ k. So, with sort, complexity of this solution will be $${\mathcal {O}}(n\log ^{2}n)$$.
Let's go in descending order by k pushing up value x from each u with val(u, x) = k. The total number of vertices visited in pushes does not exceed n·logn because dpk(u) ≤ logk(n). So, the complexity of this solution wil be $$O(n \log n)$$.