Let \(G\) be the given graph.

For each subgraph \(G'\) of \(G\), let us define a polynomial \(dp[G']\) by such that the \(k\)-th coefficient is “the number of ways to paint vertices that is impossible with \((k-1)\) operations but is possible with \(k\) operations.” For \(k=0\), we define it as \(1\), which corresponds to the only way of painting where all vertices are black. (We use polynomials for simplicity in the following explanation; it is not essential)

Every operation, except for the operation that paints a single vertex, can be interpreted as the operation of “choosing an edge \(e\) and paint all the vertices traversable from \(e\), using only the edges with weight \(e\) or less.”

When the graph is a tree and the weight of edges is pairwise distinct

Starting from a graph without an edge, consider adding the edges in the increasing order of the weight. If the edge \(e\) to be added will connect two connected components \(C_1\) and \(C_2\) to form a new connected component \(C\),

• If the operation is not performed for \(e\), the operations for \(C_1\) and \(C_2\) will be mutually independent
• If the operation is performed for \(e\), all the vertices in \(C\) will be painted red

so \(dp[C]=dp[C_1]*dp[C_2]-X^2+X\). Since this is a product of polynomials of at most degree \(K\), it can be computed in an \(O(K^2)\) time. The DP table is updated every time when an edge is added, so at a glance it seems it costs a total of \(O(NK^2)\) time to find \(dp[G]\), but in fact it is \(O(NK)\) (described below).

When there exist edges with the same weight

We will add the edges with the same weight at the same time. Suppose that connected components \(C_1,\ldots,C_m\) forms a new connected component \(C\). Just like before, by dividing into cases whether to use the new edges or not, we can show that \(dp[C]=\prod dp[C_i] - X^m+X\). Since we add \((m-1)\) edges and multiply \((m-1)\) times, the time complexity is the same as the last case.

When the graph is not a tree

We may fix a minimum spanning tree \(T\), because for each edge \(e\) that is not in \(T\), the following facts holds:

• an operation on edge included in \(T\) does not affect whether or not \(e\) exists
• an operation on edge \(e\) is the same as an operation of the operation on the edge in \(T\) with the maximum weight

Analysis

Consider connecting two connected components, depending on their number of vertices are less than \(K\) or not:

• If both of them has \(K\) or more vertices
  Such merge will occur only \(O(N/K)\) time, and each of them can be computed in an \(O(K^2)\) time, so it costs a total of \(O(NK)\) time.
• If both of them has less than \(K\) vertices Suppose that all cost required so far is considered when the size of connected component exceeds \(K\) for the first time. The time complexity of the process of connected is \(O(K^2)\), resulting in \(O(N/K)\) number of connected components, so it will cost a total of \(O(NK)\) time.
• If one has \(K\) or more vertices while the other has less than \(K\) For each element, such connection occurs at most once for the connected component with size less than \(O(NK)\), so it costs a total of \(O(NK)\) time.

Therefore, the overall time complexity is \(O(NK)\).

Reference: 木と計算量 前編 〜O(N^2)とO(NK)の木DP〜 (Trees and computational complexity, prequel - O(N^2) and O(NK) tree DP) (Blog by snuke in Japanese)