\(A\) は昇順ソートされているものとします。

重み \(A_{i}\) の辺が合計何回採用されるかが分かれば、その回数と \(A_{i}\) の積の和を計算することで答えが求まります。

クラスカル法より、最小全域木は辺の重みが小さい辺から貪欲に辺を閉路ができないなら採用すれば構築することができます。

以下の補題について考えます。

補題

\(1\) 以上 \(N\) 以下の整数 \(a\) と順列 \(P\) が与えられます。

頂点に \(1\) から \(N\) までの番号がついた頂点数 \(N\) の重み付き無向グラフ \(G\) があります。\(G\) について \(1\leq i\lt j\leq N\) を満たす任意の整数の組 \((i,j)\) に対して以下が成り立ちます。

• \(j-i\leq K\) かつ \(\max(P_{i},P_{j})\leq a\) のときのみ頂点 \(i\) と頂点 \(j\) の間に辺が存在する

\(G\) の中からいくつか辺を選びます。ただし、選ばれた辺のみからなる閉路が存在してはいけません。

選べる辺の本数の最大値を求めてください。

\(g(x)\) を全ての順列 \(P\) に対する \(a=x\) としたときの補題の答えの総和とします。

すると出力すべき値は以下のようになります。

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

よって、任意の \(1\leq i\leq N\) に対して \(g(i)\) が求まれば良いです。

\(P_{i}\leq a\) であるような添え字 \(i\) を昇順に並べた列を \(Q=(Q_{1},Q_{2},\dots Q_{a})\) としたら、この補題の答えは \(Q_{j+1}-Q_{j}\leq K\) となる \(1\leq j\lt a\) の数と等しくなります。

\(1\) 以上 \(N\) 以下からなる相異なる昇順ソートされた正整数列 \(Q=(Q_{1},Q_{2}\dots,Q_{a})\) のうち、 \(Q_{2}-Q_{1}=k\) となるようなものの場合の数は \(\left( \begin{matrix} N-k\\ a-1\\ \end{matrix} \right)\) であることから、 \(Q_{2}-Q_{1}\leq K\) となるような \(Q\) の場合の数は以下のようになります。

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

\(Q_{3}-Q_{2}\leq K,\dots ,Q_{a}-Q_{a-1}\leq K\) についても同様のことが成り立ちます。

長さ \(a\) の \(Q\) をひとつ決めると、その \(Q\) に対応する \(P\) は \(a!(N-a)!\) 通りあるので \(g(i)\) は以下のようになります。

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

階乗などを \(O(N)\) で前計算すれば、 任意の \(i\) について \(g(i)\) を \(O(K)\) で計算することができるので、全体で \(O(N(K+\log{N}))\) で答えが求まります。

実は階乗などを \(O(N)\) で前計算すれば \(g(i)\) は \(O(1)\) で求めることができるので、 \(O(N\log{N})\) でこの問題を解くこともできます (\(A\) のソートがボトルネックとなっている)。