[C++] 13278번 피보나치 합의 개수

https://www.acmicpc.net/problem/13278

---

 

26/02/14

 

 

꽤나 풀기에 까다로운 문제였다. 풀이 과정을 하나 하나씩 곱씹어보며 정리하고자 한다.

---

 

문제 접근 방식:

 

 

우리가 결국 원하는 것은 피보나치 수의 합을 $99991$로 나눈 나머지이다.($99991$은 소수다.)

따라서 피보나치 수 하나하나를 $99991$로 나누어도 상관 없다.

피사노 주기에 의하여, 피보나치 수들은 모듈로 $99991$에 대하여 어떤 주기를 가진다.

실제로 파이썬으로 계산을 해봤고, 해당 주기는 $33330$이 나왔다.

집합의 각 원소는 피보나치 수의 인덱스이므로, 집합의 각 원소들을 모두 $33330$에 대해 모듈로를 취해줄 수 있다는 사실 또한 확인할 수 있다.

이제 우리가 구하고자 하는 것은 $N$개의 집합 원소들 중 $K$개를 "중복을 제외"하고 선택했을 때, 원소의 합의 경우의 수이다.

물론 피사노 주기에 의해 $33330$미만의 합만 의미가 있으므로 중간중간 모듈로를 취해주면 된다.  

이를 생성함수 꼴로 작성하면 $\displaystyle \prod_{i=0}^{k}(1+yx^{a_{i}})$에 대해서 $y^{k}$의 계수를 구하는 것과 동일하다. ($y^{k}$의 계수는 $x$에 대한 다항식임을 생각하자.)

이에 대한 설명을 조금 덧붙이면 이렇게 설명할 수 있다.

$1$은 $a_{i}$를 선택하지 않은 경우, $x^{a_{i}}$는 $a_{i}$를 선택한 경우이고, 우리는 $k$번의 선택을 진행한다.

그때의 $y^{k}$의 계수를 본다면 $x$에 대한 다항식으로 나오는데, 해당 다항식의 특정 항의 $x$의 지수부분을 보면 $a_{i}$에 대한 합이 나옴을 확인할 수 있다.

결국 $a_{i}$들을 $k$개 뽑은 "합"을 구하고자 하는 것이므로, 위와 같은 생성함수 꼴이 나오게 된다.

하지만 문제가 있다.

여기서 $x$에 대한 다항식 자체가 최대 길이 $33330$이고(주기가 $33330$이기 때문에 그 이상의 지수는 모듈로를 취해줄 수 있음), $(1+yx^{a_{i}})$를 $N$번 곱해야 하므로 $y$에 대한 계수($x$에 대한 다항식) 또한 $N$개 저장해야한다.

하지만 문제의 제한 $N$이 최대 $50000$이라, $33330\times 50000$바이트 = 대략 $1.67GB$라 주어진 메모리로는 해당 다항식을 저장할 수 없다.

여기서 첫번째 의문이 들었다.

다시 원점으로 돌아와서 원하는 것을 생각했다.

결국 우리가 원하는 건, 합을 구하고 그 합의 인덱스에 해당하는 피보나치의 값을 곱한 값이다.

그래서 의문이 들었던 것은 $x$를 계속 유지해야 할까? 였다.

결국 $x$ 대신 어떤 것을 집어넣으면 피보나치 값을 "곱한"셈이 되지 않을까라고 생각했다.

맨 처음에 생각했던 것은 피보나치 행렬 $A = \begin{bmatrix}1 & 1\\ 1 & 0\end{bmatrix}$이였다.

그리고 나중에 $y^{k}$에 대한 계수의 합을 구한다고 하면, 행렬 제곱들로 이뤄진 것들의 합을 구하고, $1$행의 원소 두 개를 더하면 문제에서 원하는 피보나치 수의 합이 될 것이라고 생각했다.

행렬을 관리한다면 $x$에 대한 다항식의 길이 $33330$대신 숫자 $4$개로 관리할 수 있기 때문에 많은 압축을 할 수 있을 것이라고 생각했다.

여기서 두번째 의문이 들었다.

결국 $\displaystyle \prod_{i=0}^{n} (I+yA^{a_{i}})$의 $y^{k}$계수(행렬들의 합)는 어떻게 구하지?가 그 의문이었다.

$y$에 대한 다항식의 곱을 구현하려면 행렬곱을 구현해야 하고, 행렬을 계수로 가지는 다항식을 어떻게 FFT를 적용할지가 가장 큰 의문이었다. 즉, 내가 알고 있는 문제로 치환되지 않았다.

여기서 피보나치 일반항을 생각했다.

피보나치 일반항은 $\displaystyle \alpha=\frac{1+\sqrt{ 5 }}{2}, \beta=\frac{1-\sqrt{ 5 }}{2}$라고 면 $\displaystyle F(n) = \frac{\alpha^{n}-\beta^{n}}{\alpha-\beta}$이다.

내가 잘 모르는 행렬 계수 FFT 대신에, $\alpha, \beta$를 $x$에 대입해서 $\displaystyle [y^{k}]\prod_{i=0}^{k}(1+y\alpha^{a_{i}}) = \sum \alpha^{\text{idx Sum}}$랑 $\displaystyle [y^{k}]\prod_{i=0}^{k}(1+y\beta^{a_{i}}) = \sum\beta^{\text{idx Sum}}$를 계산하면 괜찮을 것이라고 생각했다.

수식적으로 정리를 해보자면, 우리가 구하고자 하는 값은 아래와 같다.
$$\sum F(\text{idx Sum}) = \frac{\sum \alpha^{\text{idx Sum}}-\sum\beta^{\text{idx Sum}}}{\alpha - \beta}$$
여기서 $\alpha, \beta$는 실수로 계산하는게 아니라 모듈로 $99991$에 대한 정수로 계산해야 한다.

즉, $x^{2}=x+1 (\text{mod }99991)$에서의 두 근 $\alpha, \beta$으로 계산하면 된다.

이러면 이제 $y$에 대한 계수가 $x$에 대한 다항식 하나가 아니라 숫자 하나로 바뀌었기 때문에 메모리도 줄어들었고, 우리가 알고있는 다항식 곱으로 여러 번 곱해서 쉽게 해결할 수 있다.

구현 시 유의 사항은, 순서대로 $K$번 곱하면 안된다는 점이다.

작은 다항식과 큰 다항식을 곱할 때, resize기준이 항상 큰 다항식 기준으로 되기 때문에 순서대로 곱하면 많이 느려지게 된다.

따라서 비슷한 크기의 다항식끼리 2개씩 묶어서 곱하면 많이 빨라지고, 세그먼트 트리마냥 트리 구조로 생각하면 이해하기 쉬울 것이다.

---

아래는 내가 위의 접근 방식과 같이 작성한 C++ 코드이다. 더보기를 누르면 확인할 수 있다.

// 13278번 피보나치 합의 개수
// FFT, 피사노 주기, 분할정복을 이용한 거듭제곱
#include <iostream>
#include <vector>
#include <cmath>

using namespace std;
#define fastio ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define endl '\n'

struct FFT {
    using ld = long double;
    using ll = long long;
    static constexpr ld PI = acosl(-1.0L);
private :
    // HELPER struct
    struct cd {
        ld re, im;
        cd(ld r=0, ld i=0): re(r), im(i) {}
        cd operator + (const cd& o) const { return cd(re+o.re, im+o.im); }
        cd operator - (const cd& o) const { return cd(re-o.re, im-o.im); }
        cd operator * (const cd& o) const { return cd(re*o.re - im*o.im, re*o.im + im*o.re); }
        cd& operator += (const cd& o) { re+=o.re; im+=o.im; return *this; }
        cd& operator -= (const cd& o) { re-=o.re; im-=o.im; return *this; }
        cd& operator *= (const cd& o) { *this = (*this)*o; return *this; }
    };
public : 
    void iterative_FTT(vector<cd>& a, int invert){
        int n = (int)a.size();
        // 1) bit-reversal permutation
        for (int i = 1, j = 0; i < n; i++){
            int bit = n >> 1;
            while (j & bit){
                j ^= bit;
                bit >>= 1;
            }
            j ^= bit;
            if (i < j) swap(a[i], a[j]);
        }
        // 2) butterflies by length = 2, 4, 8, ...
        for (int len = 2; len <= n; len <<= 1){
            ld ang = 2 * PI / len * (invert ? -1 : 1);
            cd w_len(cosl(ang), sinl(ang)); // primitive len-th root of unity on unit circle
            for (int i = 0; i < n; i += len){
                cd w(1, 0);
                for (int j = 0; j < len / 2; j++){
                    cd u = a[i+j];
                    cd v = w * a[i+j+len/2];
                    a[i+j] = u+v;
                    a[i+j+len/2] = u-v;
                    w *= w_len;
                }
            }
        }
        // 3) Divide by n (multiply by inverse) for inverse transform
        if (invert){
            ld inv_n = 1.0L / n;
            for (int i = 0; i < n; ++i){
                a[i].re *= inv_n;
                a[i].im *= inv_n;
            }
        }
    }
    // Input : Coefficient vector {a_0, a_1, ...}, {b_0, b_1, ...}
    // Output : Convolution of two coefficient vector
    vector<ll> convolution(const vector<ll>& a, const vector<ll>& b) {
        if (a.empty() || b.empty()) return {};
        int S_a = (int)a.size(), S_b = (int)b.size();
        // Make vector size easy to dnc(2^n).
        int n = 1;
        while (n < S_a + S_b - 1){
            n <<= 1;
        }
        vector<cd> fa(n), fb(n);
        for (int i = 0; i < S_a; i++) fa[i] = cd((ld)a[i], 0);
        for (int i = 0; i < S_b; i++) fb[i] = cd((ld)b[i], 0);
        // FTT
        iterative_FTT(fa, 0); iterative_FTT(fb, 0);
        // Pointwise product
        for (int i = 0; i < n; i++) fa[i] *= fb[i];
        // IFTT
        iterative_FTT(fa, 1);
        vector<ll> res(S_a + S_b - 1);
        for (int i = 0; i < res.size(); i++) {
            res[i] = (ll) llround(fa[i].re); // rounding to nearest integer
        }
        return res;
    }
};

inline int mul_mod(int a, int b, int mod) {
    return (long long)a * b % mod;
}
inline int pow_mod(int a, int e, int mod) {
    int r = 1;
    while(e > 0){
        if(e & 1) r = mul_mod(r, a, mod);
        a = mul_mod(a, a, mod);
        e >>= 1;
    }
    return r;
}
void modulo_vector(vector<long long>& a, int mod){
    for (int i = 0; i < a.size(); ++i){
        if (a[i] > mod) a[i] %= mod;
        if (a[i] < 0) a[i] += mod;
    }
    return;
}

int main(void){
    fastio

    int N, K; cin >> N >> K;
    int pisanoPeriod = 33330;
    int MOD = 99991;
    vector<int> a;
    int a_i;
    for (int i = 0; i < N; ++i){
        cin >> a_i;
        a_i %= pisanoPeriod;
        a.push_back(a_i);
    }
    int alpha = 44944, beta = 55048;
    int alpha_minus_beta_inv = 77972;
    vector<vector<long long>> y_polys_alpha, y_polys_beta;
    for (int i = 0; i < N; ++i){
        int alpha_a_i = pow_mod(alpha, a[i], MOD);
        int beta_a_i = pow_mod(beta, a[i], MOD);
        y_polys_alpha.push_back({1, alpha_a_i});
        y_polys_beta.push_back({1, beta_a_i});
    }
    FFT fft;
    while (y_polys_alpha.size() != 1){
        vector<vector<long long>> new_y_polys_alpha;
        for (int i = 0; i < (y_polys_alpha.size() / 2); ++i){
            auto res = fft.convolution(y_polys_alpha[2*i], y_polys_alpha[2*i+1]);
            modulo_vector(res, MOD);
            new_y_polys_alpha.push_back(res);
        }
        if (y_polys_alpha.size() & 1){
            new_y_polys_alpha.push_back(y_polys_alpha[y_polys_alpha.size()-1]);
        }
        y_polys_alpha = new_y_polys_alpha;
    }
    while (y_polys_beta.size() != 1){
        vector<vector<long long>> new_y_polys_beta;
        for (int i = 0; i < (y_polys_beta.size() / 2); ++i){
            auto res = fft.convolution(y_polys_beta[2*i], y_polys_beta[2*i+1]);
            modulo_vector(res, MOD);
            new_y_polys_beta.push_back(res);
        }
        if (y_polys_beta.size() & 1){
            new_y_polys_beta.push_back(y_polys_beta[y_polys_beta.size()-1]);
        }
        y_polys_beta = new_y_polys_beta;
    }
    vector<long long> sum_alpha = y_polys_alpha[0];
    vector<long long> sum_beta = y_polys_beta[0];
    int S = (sum_alpha[K] - sum_beta[K]);
    if (S < 0) S += MOD;
    int ans = mul_mod(S, alpha_minus_beta_inv, MOD);
    cout << ans;
    return 0;
}