[C++] 14958번 Rock Paper Scissors

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

---

 

26/02/04

 

 

FFT를 활용한 문자열 매칭 개념을 배울 수 있는 좋은 문제이다.

---

 

문제 접근 방식:

 

 

얼핏 보면 KMP나 라빈-카프와 같은 문자열 매칭 알고리즘을 활용하여 문제를 해결할 수 있을 것 처럼 보인다.

하지만 문제에서 요구하고자 하는 것이 실제로 매칭이 되냐 안되냐의 여부를 물어보는 것이 아니라, 몇 번 이기냐, 즉, 빈도수를 물어보기 때문에 문자열 매칭 알고리즘을 적용하는 것은 무리가 있다.

여기서 어떻게 FFT알고리즘을 적용할 수 있는지가 의문일 것이다.

먼저, 1067번 이동 문제를 풀어서 알다 시피, Convolution을 진행하는 다른쪽의 배열을 뒤집으면 일대일로 매칭할 수 있도록 만들 수 있다.

따라서 내가 고른 문자열을 뒤집어주고 시작하자.

이후, 'R', 'P', 'S'에 해당하는 문자 하나하나마다 배열을 만들어준다. 즉, 상대방의 RPS와 내 RPS, 총 6개의 배열을 만들어주면 된다.

각각의 RPS배열마다 해당 인덱스에 'R', 'P', 'S'가 존재한다면 $1$, 그렇지 않는다면 $0$으로 설정해주자.

상대방의 Rock배열과 내 Paper배열을 Convolution으로 곱하면, 상대방이 주먹을 냈을 때 내가 이기는 경우의 수를 셀 수 있다.

가위랑 보도 마찬가지 방법으로 셀 수 있다. 우리가 구하고자 하는 것은 주먹, 가위, 보 상관 없이 내가 최대한 몇 번 이길 수 있는지를 찾고 싶기 때문에, 각각의 결과를 더해주면 구해줄 수 있다.

---

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

// 14958번 Rock Paper Scissors
// NTT
/*
접근 방법:
rock, scissors, paper배열 따로 만들고 NTT 3번
*/
#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.
*/
template<int MOD, int G>
struct NTT {
    using ll = long long;
private :
    // HELPER FUNCTION
    inline int mul_mod(int a, int b) const {return (ll)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() {
    fastio

    NTT<998'244'353, 3> ntt;
    int N, M; cin >> N >> M;
    vector<int> rock(N, 0), paper(N, 0), scissors(N, 0);
    vector<int> merock(M, 0), mepaper(M, 0), mescissors(M, 0);
    char C;
    for (int i = 0; i < N; ++i){
        cin >> C;
        if (C == 'R') rock[i] = 1;
        else if (C == 'P') paper[i] = 1;
        else scissors[i] = 1;
    }
    for (int i = 0; i < M; ++i){
        cin >> C;
        if (C == 'R') merock[M-i-1] = 1;
        else if (C == 'P') mepaper[M-i-1] = 1;
        else mescissors[M-i-1] = 1;
    }
    auto paperWin = ntt.convolution(rock, mepaper);
    auto scissorsWin = ntt.convolution(paper, mescissors);
    auto rockWin = ntt.convolution(scissors, merock);
    int ans = 0;
    for (int i = M-1; i < N+M-1; ++i){
        ans = max(ans, paperWin[i]+scissorsWin[i]+rockWin[i]);
    }
    cout << ans;
    return 0;
}