1-1-1, Free Tree!

We need support vertex recolors on a weighted tree and print the sum of edge costs after every query.

An edge contributes its weight if its endpoints have different colors, otherwise $0$.

The dumb update is: repaint $v$, scan all neighbors of $v$, fix the affected edge costs. That is correct, because only edges incident to $v$ can change. But if $v$ is a star center with degree $2\cdot 10^5$, doing that for every query is very dead, very quickly.

So this is a degree-splitting problem.

The local formula

Let the old color of $v$ be $old$, and the new color be $x$.

For a neighbor $u$ connected by weight $w$:

• if $a_u = old$, then before repainting the edge was same-color and free; after repainting it becomes different-color, so the answer increases by $w$.
• if $a_u = x$, then before repainting the edge was different-color; after repainting it becomes same-color and free, so the answer decreases by $w$.
• otherwise, nothing changes.

So define:

$$S(v, col)=\sum_{(v,u)} w(v,u)\cdot [a_u=col]$$

Then, if $old \ne x$:

$$ans \mathrel{+}= S(v, old)-S(v, x)$$

If $old=x$, the delta is $0$.

That formula is the whole trick. The annoying part is answering $S(v,col)$ quickly.

Small vs heavy vertices

Pick a threshold $B$, around $450$ works fine for $2\cdot 10^5$.

A vertex is heavy if its degree is at least $B$.

For a light vertex, just scan all its neighbors to compute $S(v,old)$ and $S(v,x)$. Since its degree is small, this is cheap.

For a heavy vertex, scanning neighbors every query is too expensive. So for every heavy vertex $h$, maintain a hash map:

$$mp[h][col] = \sum_{(h,u)} w(h,u)\cdot [a_u=col]$$

Then $S(h,col)$ is just mp[h][col].

Keeping heavy maps correct

Suppose a vertex $v$ changes from $old$ to $x$.

Any heavy vertex $h$ adjacent to $v$ has one neighbor whose color changed. If the edge weight is $w$, then in $h$'s map:

• subtract $w$ from color bucket $old$
• add $w$ to color bucket $x$

So for each vertex $v$, store the list of adjacent heavy vertices, with the edge weight. After repainting $v$, update exactly those heavy maps.

This is the part people often miss. It is not enough to update the answer; future queries involving heavy vertices need their cached color sums to reflect the new colors. Cached data that lies to you is worse than no cache. Classic footgun.

Algorithm

For each test case:

1. Read colors and edges.
2. Compute the initial answer by summing weights of edges whose endpoint colors differ.
3. Choose heavy vertices by degree.
4. For every heavy vertex, build its map color -> total adjacent edge weight.
5. For every vertex, build heavyAdj[v], the adjacent heavy vertices and edge weights.
6. For each query (v, x):
   • let old = color[v]
   • compute sameOld = S(v, old) and sameNew = S(v, x):
     • if v is heavy, use its map
     • otherwise scan neighbors
   • update ans += sameOld - sameNew
   • update maps of all heavy neighbors of v
   • set color[v] = x
   • print ans

Complexity

Let $B$ be the threshold.

• Light query computation costs $O(B)$.
• Heavy query computation costs $O(1)$ average with hash maps.
• Updating heavy neighbors costs at most the number of heavy vertices adjacent to $v$, which is $O(n/B)$ in the worst case.

With $B\approx 450$ and total $n,q\le 2\cdot 10^5$, this easily fits.

Memory is $O(n + \sum \deg(h))$, where $h$ ranges over heavy vertices. Since there are few heavy vertices, this is fine.