シンプルなDPとしては、個数と和を状態に持つようなDPが考えられますが、計算量は高速化しても \(\mathrm{O}(N^4K^2)\) となり遅いので、特殊な設定を利用する必要があります。

今、平均が \(m\) であるような集合を数えたいとしましょう。

これは直線状の棒の座標 \(m\) に支点があり、座標 \(1,2,\cdots,N\) に重さ \(1\) のおもり \(K\) 個ずつまでぶら下げて釣り合うような状態の個数とみなすことができます。(\((\sum_{x \in S}x ) / |S| = m\) は\((\sum_{x \in S} (x - m) ) = 0\) と変形できるので従います)

軽実装な解法

支点より左と右に分けると、\(2\) つの集合 \(S = \{1,2,\cdots,m-1\}, T = \{1,2,\cdots,N-m\}\) の両方から何個かずつ選んで和が等しいという条件なので、dp[\(i\)][\(j\)] := \(1,2,\cdots,i\) まで考慮して和が \(j\) の集合の個数を全ての \(i, j\) について事前に求めておくと、前計算が できていれば平均 \(1\) つあたりについて \(\mathrm{O}(N^2K)\) で答える ことができます。

前計算については、前述の dp[\(i\)][\(j\)]に遷移するのは dp[\(i-1\)][\(j\)], dp[\(i-1\)][\(j - i\)]\(, \cdots, \) dp[\(i-1\)][\(j - i * K\)] なので、\(\mod i\) ごとにスライドしながら和を管理しておくと高速化でき、前計算の時間計算量は \(\mathrm{O}(N^3K)\) となります。

実装例

Formal Power Seriesを利用した解法

前計算が利用できることに気付かなくとも、\(1\) つずつ平均をずらしながら左右の和の分布を管理することが可能です (個数が一定なので、相対的な分布関係は両端しか変化しません)

前述のDPを \(\prod(1 + x^i + x^{2i} + \cdots + x^{Ki})\) を管理していると考えると、\((1 + x^i + x^{2i} + \cdots + x^{Ki}) = (1 - x^{(K+1)i})/(1-x^i)\) であるので、両端を追加・削除する操作はこの式をかけたり割ったりする操作だとみなすことができます。(戻すDP と呼ばれるテクニックと同じものです)

時間計算量は前述の解法と同じですが、空間計算量はこちらの解法の方が良いです。