【Gym - 101991K】Khoshaf（动态规划）

Gym - 101991K

题意

给定 $N,K,L,R$，构造一个长度为$N$的数列，每个数在$[L,R]$区间内，且和为3的倍数的区间恰好为$K$个，求方案数。

题解

首先转换题意，和为3的倍数即和对3取模等于0，可以先考虑每个位置放0/1/2的方案数。
然后设计状态，设 $dp[n][k][x][y][z][0/1/2]$，其中$n$是构造数列长度、$k$是和为3的倍数的区间数、$x/y/z$ 分别是和为 $0/1/2$ 的前缀数量、最后是区间总和模3的值。此时状态数有 $N^{4}K$ 那么大，显然是不可取的。通过计算可以发现，$n=x+y+z$，且 $k=[x\cdot (x+1)+y\cdot (y-1)+z\cdot (z-1)]/2$（前缀和相减得到区间和），因此前两维是不需要的。此时状态数仍然有 $N^3$ 那么多。通过上面的式子还可以发现，$x/y/z$ 并没有必要枚举到 $N$ 那么大，因为 $2\cdot k=x\cdot (x+1)+y\cdot (y-1)+z\cdot (z-1)>x\cdot (x-1)>(x-1{)}^2$，即 $x<\sqrt{2\cdot k}+1<143$，$y/z$同理。
之后考虑转移。首先求出$[L,R]$区间内对3取模等于 $0/1/2$ 的数字的数量，记为 $c[0/1/2]$。这样就可以写出转移方程：
$$dp[x][y][z][0]=\sum _{w=0}^{2}dp[x-1][y][z][w]\cdot c[(0-w)\%3]$$
$$dp[x][y][z][1]=\sum _{w=0}^{2}dp[x][y-1][z][w]\cdot c[(1-w)\%3]$$
$$dp[x][y][z][2]=\sum _{w=0}^{2}dp[x][y][z-1][w]\cdot c[(2-w)\%3]$$

代码：

#include
using namespace std;
typedef long long ll;
const int M=1e9+7;

ll d[160][160][160][3];
int main(){
	ifstream fin("khoshaf.in");
	int T,n,k,l,r;
	fin>>T;
	while(T--){
		fin>>n>>k>>l>>r;
		if(n>400){
			cout<<"0\n";
			continue;
		}
		int c[3]={(r-l+1)/3,(r-l+1)/3,(r-l+1)/3};
		if(l%3==r%3)c[l%3]++;
		if((l+1)%3==r%3)c[l%3]++,c[r%3]++;
		
		ll ans=0;
		memset(d,0,sizeof d);
		d[0][0][0][0]=1;
		for(int x=0;x<150;x++)if(x<=n){
			for(int y=0;y<150;y++)if(x+y<=n){
				for(int z=0;z<150;z++)if(x+y+z<=n){
					for(int w=0;w<3;w++){
						auto &dx=d[x][y][z][w];
						dx%=M;
						d[x+1][y][z][0]+=dx*c[(3-w)%3];
						d[x][y+1][z][1]+=dx*c[(4-w)%3];
						d[x][y][z+1][2]+=dx*c[(5-w)%3];
						if(x+y+z==n && x*(x+1)+y*(y-1)+z*(z-1)==k*2)
							ans+=dx;
					}
				}
			}
		}
		cout<<ans%M<<'\n';
	}
	return 0;
}