We assume that \(A\) is sorted in ascending order.

If we know how many times an edge with weight \(A_{i}\) is adopted, we can find the answer as the sum of the products of those counts and \(A_{i}\).

From Kruskal’s algorithm, we know that a minimum spanning tree can be constructed by greedily adopting edges with smaller weights as long as no cycles are formed.

Let’s consider the following subproblem:

Subproblem

You are given an integer \(a\) between \(1\) and \(N\), inclusive, and a permutation \(P\).

There is a weighted undirected graph \(G\) with \(N\) vertices numbered from \(1\) to \(N\). For every pair of integers \((i,j)\) such that \(1\leq i\lt j\leq N\), the following holds for \(G\).

• An edge exists between vertices \(i\) and \(j\) if and only if \(j-i\leq K\) and \(\max(P_{i},P_{j})\leq a\).

Choose some edges from \(G\). However, there must not be a cycle formed only of the chosen edges.

Find the maximum number of edges that can be chosen.

Let \(g(x)\) be the sum of the answers to the subproblem when \(a=x\) for all permutations \(P\).

Then, the value to be printed is:

\[\sum_{i=2}^{N} A_{i}(g(i)-g(i-1)).\]

Now, we need to find \(g(i)\) for any \(1\leq i\leq N\).

If we arrange the indices \(i\) such that \(P_{i}\leq a\) in ascending order to form a sequence \(Q=(Q_{1},Q_{2},\dots Q_{a})\), the answer to this subproblem is equal to the number of \(1\leq j\lt a\) such that \(Q_{j+1}-Q_{j}\leq K\).

The number of increasing sequences \(Q=(Q_{1},Q_{2}\dots,Q_{a})\) of distinct positive integers between \(1\) and \(N\) such that \(Q_{2}-Q_{1}=k\) is \(\left( \begin{matrix} N-k\\ a-1\\ \end{matrix} \right)\). Thus, the number of sequences \(Q\) such that \(Q_{2}-Q_{1}\leq K\) is:

\[\sum_{k=1}^{K} \left( \begin{matrix} N-k\\ a-1\\ \end{matrix} \right).\]

The same holds true for \(Q_{3}-Q_{2}\leq K,\dots ,Q_{a}-Q_{a-1}\leq K\).

Once we fix a \(Q\) of length \(a\), there are \(a!(N-a)!\) corresponding \(P\)’s, so \(g(i)\) is as follows:

\[g(i)=i!\cdot (N-i)!\cdot (i-1)\sum_{k=1}^{K} \left( \begin{matrix} N-k\\ i-1\\ \end{matrix} \right).\]

If we precompute factorials and such in \(O(N)\), we can calculate \(g(i)\) for each \(i\) in \(O(K)\), so the answer can be found in \(O(N(K+\log{N}))\) overall.

In fact, if factorials and such are precomputed in \(O(N)\), \(g(i)\) can be computed in \(O(1)\), so the problem can be solved in \(O(N\log{N})\) (where sorting \(A\) is the bottleneck).