A naive DP would have the count of each number and the sum as the states, but its time complexity would be \(\mathrm{O}(N^4K^2)\) even after optimization, which is too slow, so we need to exploit the particular setting.

Let us say we want to count sets whose average is \(m\).

We can consider it as the number of balanced states where there is a straight bar that pivots on a fulcrum at coordinate \(m\), and we attach at most \(K\) objects of weight \(1\) at coordinates \(1, 2, \cdots, N\). (This is because we can transform \((\sum_{x \in S}x ) / |S| = m\) to \((\sum_{x \in S} (x - m) ) = 0\).)

An easy-to-implement solution

Let us divide the bar at the fulcrum into the left and right parts. Now we have to choose two subsets from \(S = \{1,2,\cdots,m-1\}, T = \{1,2,\cdots,N-m\}\) so that the sums are equal. Thus, if we precompute dp[\(i\)][\(j\)] := (the number of sets totaling \(j\) considering \(1,2,\cdots,i\)) for every \(i\) and \(j\), we can process the queries in \(\mathrm{O}(N^2K)\) time on average per query. For the precomputation, it is dp[\(i-1\)][\(j\)], dp[\(i-1\)][\(j - i\)]\(, \cdots, \) dp[\(i-1\)][\(j - i * K\)] that transition to dp[\(i\)][\(j\)] above, so we can optimize it by maintaining the sum while “sliding” within each \(\mod i\), and the time complexity will be \(\mathrm{O}(N^3K)\).

Sample Implementation

A solution using Formal Power Series

Even if we do not notice the use of precomputation, we can still maintain the distribution of sums in the left and right parts while shifting the average one by one. (Since the counts are fixed, the relative distribution only changes at both ends.)

If we think of the DP above as maintaining \(\prod(1 + x^i + x^{2i} + \cdots + x^{Ki})\), since \((1 + x^i + x^{2i} + \cdots + x^{Ki}) = (1 - x^{(K+1)i})/(1-x^i)\), the addition and removal at both ends can be regarded as multiplying or dividing this formula.

This solution has the same time complexity as the solution above and an improved space complexity.