P3676 A Simple and Neat Data Structure Problem

Time limit: 2 s, memory limit: 256 MB.

Long long ago, there was a tree with $n$ nodes, and each node had a weight. There are $q$ operations. Each operation either modifies the weight of a node or, given a node, asks for the sum of squares of the subtree weight sums when this node is taken as the root. (If this statement is not very clear, please refer to the sample explanation.)

The first line contains two integers $n, q$. The next $n - 1$ lines each contain two integers $a$ and $b$, indicating there is an edge between $a$ and $b$ in the tree. It is guaranteed that no edges are duplicated. The next line contains $n$ integers, where the $i$-th integer is the weight of node $i$. The next $q$ lines each contain two or three numbers. If the first number is $1$, then there are two more numbers $x$ and $y$, meaning the weight of node $x$ is modified to $y$. If the first number is $2$, then there is one more number $x$, meaning you should query the sum of squares of the subtree weight sums when $x$ is the root.

For each query, output the answer.

Sample explanation: This is a star graph, with edges between $2$ and $1, 3, 4$. Initially, the node weights are $4, 3, 2, 1$. For the first query with root $2$, the subtrees of $1, 3, 4$ each contain only themselves, and the subtree of $2$ contains all nodes. Therefore, the subtree weight sums of $1, 3, 4$ are their own weights $4, 2, 1$, and the subtree weight sum of $2$ is $4 + 3 + 2 + 1 = 10$. Thus $4^2 + 2^2 + 1^2 + 10^2 = 121$. Next, modify the weight of node $1$ to $3$, so the node weights become $3, 3, 2, 1$. For the second query with root $3$, the subtrees of $1$ and $4$ each contain only themselves; the subtree of $2$ contains $1, 2, 4$; and the subtree of $3$ contains all nodes. The subtree weight sums are $3, 1, 7, 9$, so $3^2 + 1^2 + 7^2 + 9^2 = 140$. Next, modify the weight of node $2$ to $4$, so the node weights become $3, 4, 2, 1$. For the third query with root $4$, the subtree weight sums of $1$ and $3$ are $3$ and $2$, the subtree weight sum of $2$ is $3 + 4 + 2 = 9$, and the subtree weight sum of $4$ is $3 + 4 + 2 + 1 = 10$. Thus $3^2 + 2^2 + 9^2 + 10^2 = 194$. Constraints: - For 10% of the testdata, $n, q \leq 2000$. - For 40% of the testdata, $n, q \leq 6 \times 10^4$. - For another 30% of the testdata, the root in every query is $1$. - For 100% of the testdata, $1 \leq n, q \leq 2 \times 10^5$, and each input node weight satisfies $-10 \leq \text{weight} \leq 10$. It is recommended to use fast I/O, although the official solution without fast I/O can also pass. Translated by ChatGPT 5