BOJ 28342 N-beat

문제 링크입니다.

문제 요약

길이가 $B$이고 각 원소가 $0$ 이상 $2^{xy}$ 미만이며, 다음 세 조건을 만족하는 수열의 개수를 구하세요.

1. $i$번째 원소의 이진 표기법 상 $1$인 비트의 개수를 $a_i$라고 합시다.
2. $a_i \le p_1$ $(1 \le i \le B)$
3. $a_{i-1} + a_{i} \le p_2 \ (2 \le i \le B)$
4. $a_{i-2} + a_{i-1} + a_{i} \le p_3 \ (3 \le i \le B)$

풀이

우선 각 원소에 특정 개수의 비트를 할당하는 방법의 수는 이항 계수로 쉽게 구할 수 있습니다. 이를 미리 전처리해 두면, 점화식만 알면 문제를 해결할 수 있습니다. 그냥 벌리컴프-메시 짜서 점화식을 찾아도 되긴 합니다. 그게 실행 속도도 훨씬 빠릅니다…만! 그래도 직접 점화식을 구해봅시다.

직전 두 항만 참조해도 모든 조건을 만족시킬 수 있는 방법의 수를 계산할 수 있으니, 우리는 $(a_{i-2},a_{i-1})$들에 대한 항으로부터 $(a_{i-1}, a_{i})$에 대한 항을 뽑아내는 점화식을 세워야 합니다.

$p_1 \le 10$이므로 해당 순서쌍들로 행렬을 만들면 그 크기는 $121 \times 121$ 이하입니다. 이를 바탕으로 $p_2$와 $p_3$에 관한 조건을 만족시키는 경우에만 다음 항에 값을 전파하도록 행렬로 점화식을 표현하면, 행렬 거듭제곱으로 $B$번째 항을 구해낼 수 있습니다.

두 정사각행렬을 곱할 때 ${\cal O}(N^3)$의 시간이 필요하기 때문에, 테스트 케이스 당 시간 복잡도 ${\cal O}({p_1}^6 \lg B)$에 문제를 해결할 수 있습니다.

코드

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

mat operator*(mat a, mat b){
	assert(a[0].size() == b.size());
	mat c(a.size(), vector<ll>(b[0].size(), 0));
	for(int i = 0; i < c.size(); i++) for(int j = 0; j < c[0].size(); j++){
		for(int k = 0; k < a[0].size(); k++){
			c[i][j] = (c[i][j]+a[i][k]*b[k][j])%mod;
		}
	}
	return c;
}

mat id(int n){
	mat ret(n, vector<ll>(n, 0));
	for(int i = 0; i < n; i++) ret[i][i] = 1;
	return ret;
}

mat pw(mat a, ll r){
	assert(a.size() == a[0].size());
	mat ret = id(a.size());
	while(r){
		if(r&1) ret = ret*a;
		a = a*a;
		r >>= 1;
	}
	return ret;
}

ll pw(ll n, ll r){
	ll ret = 1;
	while(r){
		if(r&1) ret = ret*n%mod;
		n = n*n%mod;
		r >>= 1;
	}
	return ret;
}

void solve(){
	ll x, y, b, p1, p2, p3; cin >> x >> y >> b >> p1 >> p2 >> p3;
	mat M((p1+1)*(p1+1), vector<ll>((p1+1)*(p1+1), 0)); M[0][0] = 1;
	mat R((p1+1)*(p1+1), vector<ll>((p1+1)*(p1+1), 0));
	vector<ll> C(p1+1); C[0] = 1; for(int i = 1; i <= p1; i++) C[i] = C[i-1]*(x*y+1-i)%mod*pw(i, mod-2)%mod;
	for(int j = 0; j <= p1; j++) for(int k = 0; k <= p1; k++){
		if(j+k > p2) continue;
		for(int i = 0; i <= p1; i++){
			if(i+j > p2 || i+j+k > p3) continue;
			R[j*(p1+1)+k][i*(p1+1)+j] = (R[j*(p1+1)+k][i*(p1+1)+j]+C[k])%mod;
		}
	}
	mat ret = pw(R, b) * M;
	ll ans = 0;
	for(auto i : ret) ans = (ans+i[0])%mod;
	cout << ans << '\n';
}

int main(){
	cin.tie(0);cout.tie(0);ios::sync_with_stdio(false);
	int T; cin >> T; while(T--) solve();
}