[C++] 15050번 Laboratório de biotecnologia

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

---

 

26/02/13

 

 

누적 합을 곁들인 쉬운 FFT응용문제이다.

---

 

문제 접근 방식:

 

 

문제를 번역해서 요약을 하자면 다음과 같다.

먼저, 문자열 $S$의 가중치를 정의한다. 문자열 $S$의 가중치는, $S$를 이루고 있는 문자들의 가중치의 합으로 정의한다.

여기서 문자의 가중치란, $a=1, b=2, c=3,\dots,z=26$으로 정의한다.

우리가 구하고자 하는 것은 임의의 문자열 $S$가 주어질 때, $S$의 비어있지 않은 연속부분문자열 $s$가 가질 수 있는 서로 다른 가중치의 값이 모두 몇 개인지 구하는 것이다.

일단 연속부분문자열 $s$의 가중치는 누적 합을 통해 $\mathcal{O}(N^2)$의 시간 복잡도로 모든 값을 구할 수 있다.

하지만 여기서 $N$의 제한이 $10^{5}$으로 꽤나 크다는 사실을 발견할 수 있다. $\mathcal{O}(N^2)$으로 해결하기에는 무리가 있는 시간제한이다.

누적 합을 통해 $s$의 가중치를 구한다는 것은 배열에서 두 수 $P, Q$를 선택하여 $P-Q$를 구하는 문제로 바뀌게 된다.

이런 문제는 우리가 FFT로 많이 했던 짓이다.(20176번 Needle 문제를 해결했을 때의 방식을 생각하자.)

음수의 덧셈이라고 생각할 수 있고, 20176번 문제를 해결했을 때처럼 Offset을 둬서 해결할 수 있다. 나의 경우 Offset을 3백만 정도로 설정했다.

한 가지 구현 사항에 유의했었던 점은 평소에 NTT 모듈로 값으로 $998\ 244\ 353$을 사용했는데, 이 문제의 경우 $469\ 762\ 049$을 사용했다.

그 이유는 해당 수는 $2^{26}$의 길이를 가진 다항식까지 다룰 수 있었기 때문이다.(Offset때문에 배열의 크기가 커져서 해당 수를 사용할 수 밖에 없었다.)

---

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

// 15050번 Laboratório de biotecnologia
// FFT
/*
접근 방법:
누적 합 때려 -> P - Q where P > Q인 경우의 수 다 찾기...
두 수의 합을 구하는 경우의 수 모두 구하기 문제로 치환됨.
음수니까 빼기 대신 offset둬서 더하기로 바꾸기
*/
#include <iostream>
#include <vector>

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

/*
Notes:
MOD = 998'244'353 = 119*2^23 + 1
MOD = 1'004'535'809 = 479*2^21 + 1 / G = 3 / Works for lengths up to 2^21. / Use with CRT and 998'244'353
MOD = 469'762'049 = 7*2^26 + 1 / G = 3 / Works for lengths up to 2^26.
MOD = 167'772'161 = 5*2^25 + 1 / G = 3 / Works for lengths up to 2^25.
MOD = 1'224'736'769 = 73*2^24 + 1 / G = 3 / Works for lengths up to 2^24.
Combine above 3-NTT primes with CRT -> 75bits fast product ok.
*/
// NOTE THAT FOR FAST WORKING, CHANGE i128 -> int, int -> int
template<int MOD, int G>
struct NTT {
private :
    // HELPER FUNCTION
    inline int mul_mod(int a, int b) const {return (long long)a * b % MOD;}
    inline int pow_mod(int a, int e) const {
        int r = 1;
        while(e > 0){
            if(e & 1) r = mul_mod(r, a);
            a = mul_mod(a, a);
            e >>= 1;
        }
        return r;
    }
    inline int mod_inv(int a) const {return pow_mod(a, MOD-2);}
public : 
    void iterative_NTT(vector<int>& 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){
            // primitive len-th root of unity
            int w_len = pow_mod(G, (MOD - 1) / len);
            if (invert) w_len = mod_inv(w_len);
            for (int i = 0; i < n; i += len){
                int w = 1;
                for (int j = 0; j < len / 2; j++){
                    int u = a[i+j];
                    int v = mul_mod(w, a[i+j+len/2]);
                    int x = u+v; if (x >= MOD) x -= MOD;
                    int y = u-v; if (y < 0) y += MOD;
                    a[i+j] = x;
                    a[i+j+len/2] = y;
                    w = mul_mod(w, w_len);
                }
            }
        }
        // 3) Divide by n (multiply by inverse) for inverse transform
        if (invert){
            int inv_n = mod_inv(n);
            for (int i = 0; i < n; ++i){
                a[i] = mul_mod(a[i], inv_n);
            }
        }
    }
    // Input : Coefficient vector {a_0, a_1, ...}, {b_0, b_1, ...}
    // Output : Convolution of two coefficient vector
    vector<int> convolution(vector<int> a, vector<int> 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;
        }
        a.resize(n); b.resize(n);
        // Normalize to [0, MOD)
        for (int i = 0; i < n; ++i){
            a[i] %= MOD; if (a[i] < 0) a[i] += MOD;
            b[i] %= MOD; if (b[i] < 0) b[i] += MOD;
        }
        // NTT
        iterative_NTT(a, 0); iterative_NTT(b, 0);
        // Pointwise product
        for (int i = 0; i < n; i++) {
            a[i] = mul_mod(a[i], b[i]);
        }
        // INTT
        iterative_NTT(a, 1);
        a.resize(S_a + S_b - 1);
        return a;
    }
};

int main(void){
    fastio

    string S; cin >> S;
    vector<int> A(6'000'000, 0), B(6'000'000, 0);
    int R = 0;
    A[R+3'000'000] = 1; B[-R+3'000'000] = 1;
    for (int i = 0; i < S.size(); ++i){
        R += ((int)(S[i]-'a')+1);
        A[R+3'000'000] = 1; B[-R+3'000'000] = 1;
    }
    NTT<469'762'049, 3> ntt;
    auto C = ntt.convolution(A, B);
    int ans = 0;
    for (int i = 6'000'000+1; i < C.size(); ++i){
        if (C[i] > 0) ++ans;
    }
    cout << ans;
    return 0;
}