【CodeChef】【DP】Count Subsequences

CodeChef CSUBSQ Count Subsequences

题目大意

◇题目传送门◆

如果一个整数序列的各元素之和可以被给定整数 $K$ 整数，则称这个序列是好的。给定整数序列 $A_1,A_2,\dots ,A_N$。Hasan 想计算该序列有多少子序列是好的。不过这样的话问题就太简单了，因此 Hasan 决定再加一个限制。

任意非空下标序列 $i_1<i_2<\cdots <i_L$ 对应子序列 $A_{i_1},A_{i_2},\dots ,A_{i_L}$，如果满足 $i_L-i_1\ge W$ 则称这个子序列非常好。请帮 Hasan 统计非常好的子序列方案数。由于答案可能非常大，请输出方案数对 $10^9+7$ 取模的结果。

分析

假设没有 $W$ 这个限制，那么这个问题就变得简单了，就是一个简单的背包问题。

我们如果先固定左端点 $L$，那么我们需要的就是在区间 $[L+W,N]$ 中查询一个元素，使它与左端点 $L$ 之间选出任意多个元素，它们的和模 $K$ 等于 $0$。

那么可以设计一个这样的状态：

状态 $f(L,i,j)$ 表示当前区间左端点为 $L$ （外层循环枚举），当前加入第 $i$ 个数字，和模 $K$ 为 $j$ 的方案数。

不难得出以下初始状态和转移：

初始状态：$f(L,L,0)=1$ 和 $\forall j\in [1,K-1],f(L,L,j)=0$。

转移：$f(L,i,j)=f(L,i-1,j)+f(L,i-1,(j-A_i+K)\mathrm{m}\mathrm{o}\mathrm{d}K)$。

显然最后的答案就是对于每个左端点 $L$，在区间 $[L+1,N]$ 中选择元素之和模 $K$ 等于 $0$ 的方案数 减掉 在区间 $[L+1,L+W-1]$ 中选择元素之和模 $K$ 等于 $0$ 的方案数。

显然直接采用这种方法枚举左端点 $L$ 计算就可以超时。

那么考虑采用分治的方法进行优化。

设当前我们正在处理的区间为 $[L,R]$，我们将它拆分成两个子区间 $[L,mid]$ 和 $[mid+1,R]$。（$mid$ 是当前区间的中点），通过递归得到这两个子区间的答案。

那么我们就定义两个 DP 状态：$f_L(p,k),f_R(p,k)$，分别表示在区间 $[p,mid]$ 上有多少个序列的和模 $K$ 等于 $k$，区间 $[mid,p]$ 上有多少个子序列的和模 $K$ 等于 $k$。

计算这两个 DP 状态的方法和上面介绍的一样。

那么考虑如何计算当前点的贡献。

我们记当前左边的元素之和模 $K$ 等于 $k_l$，右边的元素之和模 $K$ 等于 $k_r$。需要保证 $k_l+k_r$ 模 $K$ 等于 $0$。

那么我们就先枚举 $[L,mid]$ 的一个左端点 $l$，强制它要选（也就是它就是第一个）。那么如何计算这个如何计算，我们分成两种情况讨论：

---

情况1： 当我们所选取的 $l+W-1\ge mid$ 时。

这时的答案就是必须选 $l$，在区间 $[l+1,mid]$ 里任意选择的方案数乘上在区间 $[mid+1,l+W-1]$ 随意选且在区间 $[l+W,R]$ 选择至少一个的方案数。

用我们计算出的 DP 值表示就是：

$$(f_L(l,k_l)-f_L(l+1,k_l))\times (f_R(R,k_r)-f_R(l+W-1,k_r))$$

情况2： 当我们所选取的 $l+W-1<mid$ 时。

我们发现无法直接计算出相应的答案，但我们可以将它拆成两部分：

1. 必须选择 $l$，且在 $[mid+1,R]$ 中选至少一个，在 $[l+1,mid]$ 中随便选的方案数：可以利用我们计算的 DP 表示为 $(f_L(l,k_l)-f_L(l+1,k_l))\times f_R(R,k_r)-f_R(mid,k_r)$；
2. 必须选择 $l$，且在 $[mid+1,R]$ 中不选，在 $[l+W,mid]$ 中选至少一个，在 $[l+1,l+W-1]$ 中随便选的方案数，这个我们采用递归的方法进行处理。

---

总时间复杂度为 $O(NK{\log }_{2}N)$。

参考代码

#include <cstdio>
#include <algorithm>
using namespace std;

typedef long long ll;
const int Maxn = 1e5;
const int Maxk = 50;
const ll Mod = 1e9 + 7;

int N, K, W;
int A[Maxn + 5];

ll fLef[Maxn + 5][Maxk + 5], fRig[Maxn + 5][Maxk + 5];
ll ans;

void Divide(int lb, int ub) {
	if(ub - lb < W) return;
	if(lb == ub) {
		if(A[lb] == 0) ans = (ans + 1) % Mod;
		return;
	}
	int mid = (lb + ub) >> 1;
	Divide(lb, mid), Divide(mid + 1, ub);
	for(int k = 0; k < K; k++)
		fRig[mid][k] = fLef[mid + 1][k] = 0;
	fRig[mid][0] = 1;
	for(int j = mid + 1; j <= ub; j++)
		for(int k = 0; k < K; k++)
			fRig[j][k] = (fRig[j - 1][k] +
				fRig[j - 1][(k - A[j] + K) % K]) % Mod;
	fLef[mid + 1][0] = 1;
	for(int j = mid; j >= lb; j--)
		for(int k = 0; k < K; k++)
			fLef[j][k] = (fLef[j + 1][k] +
				fLef[j + 1][(k - A[j] + K) % K]) % Mod;
	ll tot = 0;
	for(int lefsum = 0; lefsum < K; lefsum++) {
		int rigsum = (K - lefsum + K) % K;
		for(int j = lb; j + W <= ub && j <= mid; j++) {
			ll lef = (fLef[j][lefsum] - fLef[j + 1][lefsum] + Mod) % Mod, rig;
			if(j + W - 1 >= mid)
				rig = (fRig[ub][rigsum] - fRig[j + W - 1][rigsum] + Mod) % Mod;
			else rig = (fRig[ub][rigsum] - fRig[mid][rigsum] + Mod) % Mod;
			tot = (tot + lef * rig % Mod) % Mod;
		}
	}
	ans = (ans + tot) % Mod;
}

int main() {
#ifdef LOACL
	freopen("in.txt", "r", stdin);
	freopen("out.txt", "w", stdout);
#endif
	int _;
	scanf("%d", &_);
	while(_--) {
		scanf("%d %d %d", &N, &K, &W);
		for(int i = 1; i <= N; i++) scanf("%d", &A[i]);
		ans = 0;
		Divide(1, N);
		printf("%lld\n", ans);
	}
	return 0;
}