【CodeChef】【DP】Count Subsequences

CodeChef CSUBSQ Count Subsequences

題目大意

◇題目傳送門◆

如果一個整數序列的各元素之和可以被給定整數 $K$ 整數，則稱這個序列是好的。給定整數序列 $A_1, A_2, \ldots, A_N$。Hasan 想計算該序列有多少子序列是好的。不過這樣的話問題就太簡單了，因此 Hasan 決定再加一個限制。

任意非空下標序列 $i_1< i_2 < \cdots < i_L$ 對應子序列 $A_{i_1}, A_{i_2}, \ldots , 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) \mod 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;
}