I afraid I don’t understand the problem statement. What do you mean by “in each iteration” and “smallest edge”?

Well, we can keep DSU in persistent array structure (it can be implemented as persistent segment tree, or as history of changes for all elements of the array) and do binary search over it. That is, in DSU we need to maintain not only size of connected component but also sum of a[i] over vertices.

Keeping track of the sum of $$$A$$$ along with the connected components is a classic problem that I'm sure you know the solution of. We will iterate through the edges and join them using DSU. To answer the queries with the information we're keeping track of using DSU, we will maintain a set of query indices for each vertex (so insert $$$i$$$ into the sets of the verticies $$$U_i$$$ and $$$V_i$$$). While merging the sets, we can easily check whether any $$$i$$$ exists so that both components' sets include $$$i$$$. When we detect such an event, we can then easily get the sum of $$$A$$$ in both connected components.

Final runtime $$$O(N \,log \,N\, log\, N)$$$.

You could do smth like this. We will keep a priority queue of all edges that we can currently use and each turn we will take the one we the smallest index and unite components with this edge. For every connected component keep a set (sorted by w_i) of all the edges that are connected to any vertex in the component, and also current sum of a's. When you unite two components, you can merge these two sets using small-to-large. After you merged them and updated the sum of a's of our new component, you can find all edges that now have their w_i less than a sum of a's (it will be some prefix of edges from the set). So, for every edge that has its w_i less than a sum, you delete it from the set and increase its counter by one. When the counter for an edge reaches 2, it means that now you can use it so you just push it to the priority queue.

It works in O(nlog^2)

I assume that $$$a[i], w \geq 0$$$ as this is quite common and I have no idea how to solve the problem without this constraint. With this constraint however, I think that I can solve the problem in $$$O(N \log(N) \log(A))$$$ where $$$A$$$ is the maximum value of $$$a[i]$$$ and $$$w$$$:

We say that an edge is good if the sum of its two components is at least $$$w$$$. When the edge with index $$$i$$$ becomes good, we insert it into a priority queue. We repeatedly pop the smallest $$$i$$$, check if $$$u_i$$$ and $$$v_i$$$ are still in different components and merge them if they are. Now, the hard part of the problem is How to find the edges which become good when we merge two components?

Let $$$x_i$$$ denote $$$w_i - ($$$ sum of components of $$$i$$$-th edge $$$)$$$. So initially $$$x_i = w_i - a[u_i] - a[v_i]$$$. Whenever we merge the component of $$$u_i$$$ with some other component of sum $$$s$$$, we need to decrease $$$x_i$$$ by $$$s$$$. Edge $$$i$$$ becomes good, when $$$x_i$$$ becomes $$$\leq 0$$$. Of course, we can't just iterate over all edges which have an endpoint in the merged components — we might iterate over every edge $$$N$$$ times! So instead of applying all updates one by one, we perform them lazily:

Assume that edge $$$i$$$ has endpoints in two different components. In one component we still have to apply the lazy update $$$y_1$$$ and in the other $$$y_2$$$. This means that the correct value of $$$x_i$$$ after applying all updates is actually $$$\tilde{x_i} = x_i - y_1 - y_2$$$. If both $$$y_1, y_2 < \frac{x_i}{2}$$$, we don't need to apply those updates now, because $$$\tilde{x_i} > 0$$$ and thus edge $$$i$$$ does not become good by applying those updates. So we wait with applying updates until some $$$y_j \geq \frac{x_i}{2}$$$. This gives us $$$\tilde{x_i} \leq \frac{x_i}{2}$$$. Since $$$x_i$$$ is halved by each update, there can be at most $$$O(\log(A))$$$ such updates!

Some more details:

• We have a DSU where every component $$$c$$$ stores the sum $$$s_c$$$ of all $$$a$$$ and a priority queue which stores pairs $$$(z, i)$$$. Here $$$i$$$ is the index of an edge with an endpoint in the component and $$$z = \hat{x_i} + 2 \cdot \hat{s_c}$$$ where $$$\hat{x_i}$$$ and $$$\hat{s_c}$$$ refer to the values at the last time when $$$x_i$$$ was modified.
• The lazy update which still needs to be applied is $$$y = s_c - \hat{s_c}$$$. And we want to perform it when $$$y \geq \frac{x_i}{2} \Leftrightarrow 2 \cdot s_c \geq x_i + 2 \cdot \hat{s_c} = z$$$. So we keep popping elements from the priority queue as long as $$$z \leq 2 \cdot s_c$$$ and apply the lazy updates on the corresponding edges.
• To merge the priority queues of two components, we use small-to-large merging.

Do you have a link to the problem? I would like to test if my implementation is correct.