There are two ways this problem can be solved. I really like these solutions.
- Square Root Decomposition
Keep a set of at most O(sqrtN) "active nodes".
When you have a query of type 2, just loop through all these active nodes and add the distance between each one to the queried node. When you have a query of type 1, add the new node to the set if the size does not exceed sqrtN. When the size of this set exceeds sqrtN, we need to basically somehow empty the set and update all nodes with their total distance to all active nodes. But this can be easily done in O(N), not O(NsqrtN). The idea is that let's say the root of the tree is node 0, and first find the total distance sum of all active nodes to the root in O(N). Then, consider a node n and its child c, and let ans[n] = the total distance sum of all active nodes to node n. When we move from node n to its child c, we move one closer to every active node in the subtree of c, while moving one further from every other active node in the rest of the tree. So dp[c] = dp[n] - (sub[c]) + (total number of active nodes - sub[c]), where sub[c] is the number of active nodes in the subtree of c (and all sub[..] can be precomputed in O(N), with the same approach one would use to find subtree sizes). So this transition is O(1), meaning that we need O(N) time to simultaneously apply all updates (note that this works for any number of active nodes). In the end, the total time complexity is O(QsqrtN) (since each query of type 2 is O(sqrtN) in the worst case, and queries of type 1 are O(Q/sqrtN * N) = O(QsqrtN) amortized).
BUT: Overall, this is quite slow for me, taking 5 or 6 seconds to process 200k queries over a tree with 100k nodes.
One example of this sort of reasoning:
- Heavy Light Decomposition
Use the same recurrence as above with the same definitions: dp[c] = dp[n] + sz - 2 * sub[c], where sz = total number of active nodes in the whole tree. Now, what happens if we substitute the recurrence in for dp[n]? Then we get dp[c] = (dp[parent of n] + N - 2 * sub[n]) + N - 2 * sub[c]. Reapplying this recurrence until we reach the root (say, node 0) gets us
dp[c] = dp + (depth of node c) * sz - 2 * (for every node i on the path from c to the root, except for the root itself) sub[i].
IN OTHER WORDS WE NOW HAVE AN EPIC CLOSED-FORM EXPRESSION FOR THE ANSWER TO ANY NODE
It's clear that dp is easy to maintain (when we get a type 1 query with a new node, just add its distance from the root, which just equals its depth, to dp). sz is also trivial to maintain (it's just the total number of type 1 queries received up to a given point).
But how can we efficiently find 2 * (for every node i on the path from c to the root, except for the root itself) sub[i]?
Heavy Light Decomposition. Basically we need to support two types of queries:
- Update every value on a path to the root by 1
- Find the sum of all values on a path from the root (and then subtract out sub[root] since remember we don't count that in the expression)
And this can be done with Heavy-Light Decomposition incorporating a Lazy Segment Tree.
For the same bounds as above (200k queries on a tree of 100k nodes), I had a runtime of less than 300ms.