Set partition refinement with subtrees
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.