T765905 计数(count) 赛后题解

题目传送门

思路

为了方便，我们称原题中的 \(n,k\) 为 \(N,K\)。

发现原题可以转化为：求所有的数组 \(a_1,a_2\dots a_m\)，使得

• 对于所有 \(1\le i<m\)，满足 \(a_i+a_{i+1}<K\)
• \(\sum\limits_{i=1}^{m-1}(a_i+1)(a_{i+1}+1)=N\)

做法

考虑 dp。
定义 dp[i][j] 为和为 i，且当前一位为 j 的方案数。
最开始假设 \(a_0\) 为 \(0\sim K-1\)，即 dp[0][0~K-1]=1。
答案为 \(\sum\limits_{i=0}^{K-1}dp[n][i]\)。
转移时枚举当前的和 i、当前要选的数 j、上一次要选的数 k。如果满足 \((j-1)(k-1)\le i\) 且 \(j+k<K\)，则作转移。

如果 \((j-1)(k-1)>i\)，显然可以跳过剩下的对 k 的枚举。
所以时间复杂度：\(O(Tn^2\log n)\)，可以通过 \(1\le T\le10\)，\(1\le k\le n\le2500\)。

代码

#include<bits/stdc++.h>
using namespace std;
const int mod = 998244353;
int dp[2505][2505];
signed main()
{
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
	freopen("count.in", "r", stdin);
	freopen("count.out", "w", stdout);
	int _;
	cin >> _;
	while(_ --)
	{
		int n, K;
		cin >> n >> K;
		memset(dp, 0, sizeof(dp));
		for(int i = 0; i < K; i ++) dp[0][i] = 1;
		for(int i = 1; i <= n; i ++)
		{
			for(int j = 0; j < K; j ++)
			{
				for(int k = 0; k < K; k ++)
				{
					if(i - (long long)(j + 1) * (k + 1) < 0 || j + k >= K) break;
					dp[i][j] += dp[i - (long long)(j + 1) * (k + 1)][k];
					if(dp[i][j] >= mod) dp[i][j] -= mod;
				}
			}
		}
		int ans = 0;
		for(int i = 0; i < K; i ++) ans = (ans + dp[n][i]) % mod;
		cout << ans << endl;
	}
	return 0;
}