BOJ 1336 수열의 개수 NKD

문제 링크입니다.

거품 낀 제한

길이가 $K$고 원소의 합이 최소인 수열 $A$는 $[1, 2, 3, \cdots , K]$ 임이 자명하고, 이때 원소의 합은 $K(K+1)/2$입니다. $N \ge K(K+1)/2$여야 하는데, 제한에 의하면 $K$가 500만 돼도 $N$이 $100\,000$을 넘깁니다. $K$ 제한을 500인 것처럼 생각합시다.

계단 쌓기

수열 [1, 2, 3, …, $K]$를 가져옵시다. 이 수열을 계단처럼 생각해, 계단을 규칙에 맞게 변형하는 방법을 생각해 봅시다.

계단의 첫 칸부터 시작해, 순서대로 높이를 결정합니다. 현재 상황에서는 최소 0에서 최대 $D-1$까지 높이를 올릴 수 있습니다.

계단이 증가해야 한다는 조건이 있습니다. 어떤 칸을 올릴 때, 이 칸과 그 뒤의 칸들은 모두 높이 차가 1인 계단을 이루고 있기 때문에, 계단 모양이 깨지지 않기 위해서는 그 뒤의 칸까지 모두 올려야 합니다.

따라서, 1번째 칸을 1 올리면 수열의 총합은 $K$ 증가하고, 2번째 칸을 1 올리면 수열의 총합은 $K-1$ 증가하고 … 이런 식입니다.

DP

이제 문제를 1, 2, …, $K$를 각각 $D-1$개 이하로 사용해 그 합을 $N-K(K+1)/2$로 만드는 방법의 수를 구하는 것으로 변형시켰습니다. 1부터 $K$까지 각 수 $X$마다, 작은 $i$부터 보면서 $i$를 만드는 방법의 수를 $i, i-X, i-2X,…,i-(D-1)X$를 만드는 방법의 수의 합으로 계산할 수 있습니다.

이를 위해 $X$씩 점프하는 누적 합을 관리하면 각 $X$마다 $\mathcal{O}(N)$에 모든 갱신을 완료할 수 있습니다.

$X$는 아무리 커봤자 500이니 제한 시간 안에 문제를 해결할 수 있습니다. 정확히는, $K(K+1)/2 \le N$으로부터 $K$의 상한이 $\mathcal{O}(\sqrt{N})$임을 알 수 있습니다. 따라서 시간 복잡도는 $\mathcal{O}(N\sqrt{N})$입니다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pll = pair<ll, ll>;
const ll mod = 1e9+7;

ll n, k, m;

ll d[101010], psum[101010];

int main(){
	cin.tie(0); cout.tie(0); ios::sync_with_stdio(false);
	cin >> n >> k >> m;
	if(k*(k+1)/2 > n) return !(cout << 0);
	if(k*(k+1)/2*m < n) return !(cout << 0);
	n -= k*(k+1)/2;
	d[0] = psum[0] = 1;
	for(ll x = 1; x <= k; x++){
		for(int i = 1; i <= n; i++){
			psum[i] = d[i];
			if(i >= x) psum[i] += psum[i-x];
			psum[i] %= mod;
			d[i] = psum[i];
			if(i >= m*x) d[i] -= psum[i-m*x];
			d[i] += 2*mod;
			d[i] %= mod;
		}
	}
	cout << d[n];
}