Set partition refinement with subtrees

Question

In trying to design an algorithm, I needed a datastructure to implement a restricted kind of set partition refinement, where the sets $X$ to split on are subtrees. Specifically, given an arbitrary tree $T$ (computed earlier in the algorithm), and an initial partition where everything is in the same set, I need to do two things efficiently:

1. Refine the partition by splitting nodes in some subtree $T'$ of $T$ from the rest of the nodes.
2. Given a node $n$ in the tree, figure out the farthest (closest to the root) ancestor in its subset $S_n$.

The algorithm may have to do both up on (close to) every node on the tree, in any order. There's an obvious algorithm, similar to the one described on Wikipedia, where each node stores a pointer to the farthest ancestor in its set, and then when given a subtree $T'$ to split, we walk through every node in $T'$, and update the pointer to the root $r$ of $T'$ if it's an ancestor of $r$. The problem with this implementation is that since splitting can take up to $O(n)$ time (where $n$ is the size of the tree) to update all the pointers, the whole thing would be $O(n^2)$. Another implementation would be store, for each node, whether it's been split with its parent node. This makes splitting $O(1)$, since we just have to set the flag for the root node, but now finding the ancestor is $O(n)$ (we have to walk up the tree) so the whole thing is still $O(n^2)$. Is there a way to make this faster than $O(n^2)$ by taking advantage of the structure of the problem?

Answer 1

Yes. You can build a data structure with the property that each operation takes at most $O(\log n)$ time. Thus, your overall running time for doing $n$ operations will be $O(n \log n)$ instead of $O(n^2)$ -- a significant improvement.

As a warmup, imagine that the tree $T$ was balanced and in particular had depth $O(\log n)$. Then it is immediate to see how to achieve $O(\log n)$ time per operation; your second approach suffices. In particular, you set a flag in each node that is a split point. Since the tree has height $O(\log n)$, walking up the tree takes $O(\log n)$ time.

What if $T$ isn't balanced? We can generalize that answer. Number the leaves of $T$ with the numbers $1,2,\dots,n$, from left to right. Now build a new balanced binary tree $T^*$ with leaves $1,2,\dots,n$. The shape of $T^*$ will be unrelated to the shape of $T$, but by construction the depth of $T^*$ will be $O(\log n)$.

Now, we store flags inside $T^*$ instead of inside $T$. In particular, we will store a flag inside a node $v \in T^*$ if (1) the entire subtree of $v$ (including $v$ itself) is in the same partition, and (2) the parent of $v$ is in a different partition from $v$. The flag identifies the partition that $v$ is in. To find the partition that a node is in, you just walk up $T^*$, which takes $O(\log n)$ time.

And here is the key part. When you split a partition, you can update the flags in $O(\log n)$ time. In particular, when you split off a subtree of $T$, that corresponds to a set of leaves that are in a consecutive interval (namely, leaves numbered $i,i+1,i+2,\dots,j$ for some $i,j$). But over in $T^*$, every consecutive interval is covered by a (disjoint) union of at most $O(\log n)$ subtrees of $T^*$. So, you update the flags of the roots of those subtrees to reflect the new partition they are in. That requires only $O(\log n)$ updates, so can be done in $O(\log n)$ time.