BOJ - 19165번 Addition Robot

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

19165번: Addition Robot

ICPC 2019 Jakarta Regional에 나온 문제로 구현이 상당히 귀찮았지만(출제자의 의도와 맞는 풀이인지는 모르겠지만) 아이디어가 재밌는 문제 였습니다.

 

 

먼저 "A"와 "B"로만 이루어진 특정 문자열이 주어집니다.

그 후 특정한 로봇이 해당 문자열을 가지고 2가지 행동을 한다고 합니다.

1. 특정 구간 L ~ R 사이의 문자를 전부 Toggle 합니다.
   이때 Toggle이란 "A"라는 문자를 "B"로, "B"라는 문자는 "A"로 바꾸는 것을 의미합니다.
2. 특정 구간 L ~ R과 A와 B라는 임의의 숫자가 주어질때 아래 함수 f와 같이 행동합니다.

function f(L, R, A, B):
  FOR i from L to R
    if S[i] = 'A'
      A = A + B
    else
      B = A + B
  return (A, B)

 

이때 2번 행동을 할때마다 결과로 나오는 (A,B)의 결과를 출력하면 되는 문제입니다.

먼저 구간의 데이터를 전부 Toggle한다는 것으로부터 구간의 데이터에 특정 연산을 가할때 매우 효율적이게 할 수 있는 자료구조가 필요하다는 것을 알 수 있습니다.

이러한 자료구조는 lazy propagation segment tree가 있습니다.

또한 Toggle은 A → B, B → A로 가는 것으로 부터 flip 연산과 같다는 것을 파악할 수 있습니다.

 

 

그럼 이제 1번 행동에 대한 통찰은 끝났으니 2번 행동에 대해 생각해 보겠습니다.

2번 행동은 i번째 데이터가 "A"인 경우 A에 B를 더해주고, i번째 데이터가 B인 경우 B에 A를 더해줍니다.

그렇다면 "AAABA"와 같은 경우는 총 5번 연산을 해줘야 하는데 이경우 한번 연산때 마다 $O(N)$의 시간복잡도를 가지며 TLE를 피할 수 없습니다.

"ABA"와 "AAB"에 대해 한번 계산을 해보겠습니다.

"ABA"에 대해 계산하기

"A": (A+B, B)

"B": (A+B, A+B+B) = (A+B, A+2B)

"A": (A+B+A+2B, A+2B) = (2A+3B, A+2B)

(2A+3B, A+2B)가 리턴됩니다.

"AAB"에 대해 계산하기

"A": (A+B, B)

"A": (A+B+B, B) = (A+2B, B)

"B": (A+2B, B+A+2B) = (A+2B, A+3B)

(A+2B, A+3B)가 리턴됩니다.

 

위 결과로부터 로봇의 행동에는 A와 B의 순서 또한 상관이 있다는 것을 알 수 있습니다.

즉, 교환 법칙이 통하지 않는다는 것이죠.

상당히 당황스럽지만 다행이게도 이러한 성질을 가지는 공간에 대해서 이미 배운적 있습니다.

바로 행렬 공간이 이러한 성질을 띄죠.

그러므로 행렬로 풀이를 해보도록 하겠습니다.

"A"를 만날 경우 A에 B를 더하므로 아래와 같이 표현 가능합니다.

반대로 "B"를 만날 경우 B에 A를 더하므로 아래와 같이 표현 가능합니다.

또한 행렬은 $CDC$와 같이 있을때 $CDC$에 대해 미리 계산해 놓으면 이후 들어오는 입력에 대해 한번의 행렬 연산으로 계산을 마칠 수 있다는 점에서 segment Tree의 연산에도 적합합니다.

위 아이디어를 통해 구현한 코드는 다음과 같습니다.

#include<iostream>
#include<algorithm>
#include<string.h>

using namespace std;
int N, Q;
const long long mod = 1e9 + 7;
pair<pair<long long, long long>, pair<long long, long long>> segTree[100004 * 4];
bool lazy[100004 * 4];

void setLazy(int ptr, int s, int e)
{
    pair<long long, long long> A = { segTree[ptr].second.second, segTree[ptr].second.first };
    pair<long long, long long> B = { segTree[ptr].first.second, segTree[ptr].first.first };
    segTree[ptr] = { A,B };
    if (s != e)
    {
        lazy[ptr * 2] = !lazy[ptr * 2];
        lazy[ptr * 2 + 1] = !lazy[ptr * 2 + 1];
    }
    lazy[ptr] = 0;
}

void init(int ptr, int s, int e, int idx, pair<pair<long long, long long>, pair<long long, long long>> val)
{
    if (s > idx || e < idx) return;
    if (s == e)
    {
        segTree[ptr] = val;
        return;
    }
    init(ptr * 2, s, (s + e) / 2, idx, val);
    init(ptr * 2 + 1, (s + e) / 2 + 1, e, idx, val);
    pair<long long, long long> A = { (segTree[ptr * 2 + 1].first.first * segTree[ptr * 2].first.first + segTree[ptr * 2 + 1].first.second * segTree[ptr * 2].second.first) % mod, (segTree[ptr * 2 + 1].first.first * segTree[ptr * 2].first.second + segTree[ptr * 2 + 1].first.second * segTree[ptr * 2].second.second) % mod };
    pair<long long, long long> B = { (segTree[ptr * 2 + 1].second.first * segTree[ptr * 2].first.first + segTree[ptr * 2 + 1].second.second * segTree[ptr * 2].second.first) % mod, (segTree[ptr * 2 + 1].second.first * segTree[ptr * 2].first.second + segTree[ptr * 2 + 1].second.second * segTree[ptr * 2].second.second) % mod };
    segTree[ptr] = { A,B };
}

void update(int ptr, int s, int e, int l, int r)
{
    if (lazy[ptr]) setLazy(ptr, s, e);
    if (s > r || e < l) return;
    if (l <= s && e <= r)
    {
        pair<long long, long long> A = { segTree[ptr].second.second, segTree[ptr].second.first };
        pair<long long, long long> B = { segTree[ptr].first.second, segTree[ptr].first.first };
        segTree[ptr] = { A,B };
        if (s != e)
        {
            lazy[ptr * 2] = !lazy[ptr * 2];
            lazy[ptr * 2 + 1] = !lazy[ptr * 2 + 1];
        }
        return;
    }
    update(ptr * 2, s, (s + e) / 2, l, r);
    update(ptr * 2 + 1, (s + e) / 2 + 1, e, l, r);
    pair<long long, long long> A = { (segTree[ptr * 2 + 1].first.first * segTree[ptr * 2].first.first + segTree[ptr * 2 + 1].first.second * segTree[ptr * 2].second.first) % mod, (segTree[ptr * 2 + 1].first.first * segTree[ptr * 2].first.second + segTree[ptr * 2 + 1].first.second * segTree[ptr * 2].second.second) % mod };
    pair<long long, long long> B = { (segTree[ptr * 2 + 1].second.first * segTree[ptr * 2].first.first + segTree[ptr * 2 + 1].second.second * segTree[ptr * 2].second.first) % mod, (segTree[ptr * 2 + 1].second.first * segTree[ptr * 2].first.second + segTree[ptr * 2 + 1].second.second * segTree[ptr * 2].second.second) % mod };
    segTree[ptr] = { A,B };

}

pair<pair<long long, long long>, pair<long long, long long>> getVal(int ptr, int s, int e, int l, int r)
{
    if (lazy[ptr]) setLazy(ptr, s, e);
    if (s > r || e < l) return { {1,0}, {0,1} };
    if (l <= s && e <= r)
    {
        return segTree[ptr];
    }
    pair<pair<long long, long long>, pair<long long, long long>> temp1 = getVal(ptr * 2, s, (s + e) / 2, l, r);
    pair<pair<long long, long long>, pair<long long, long long>> temp2 = getVal(ptr * 2 + 1, (s+e)/2+1, e, l, r);
    pair<long long, long long> A = { (temp2.first.first * temp1.first.first + temp2.first.second * temp1.second.first) % mod, (temp2.first.first * temp1.first.second + temp2.first.second * temp1.second.second) % mod };
    pair<long long, long long> B = { (temp2.second.first * temp1.first.first + temp2.second.second * temp1.second.first) % mod , (temp2.second.first * temp1.first.second + temp2.second.second * temp1.second.second) % mod };
    return { A,B };
}

int main()
{
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    cin >> N >> Q;
    for (int i = 0; i <= 4*N; i++)
        segTree[i] = { {1,0}, {0,1} };
    for (int i = 1; i <= N; i++)
    {
        char temp; cin >> temp;
        if (temp == 'A')//a인 경우 pair<pair<int,int>,pair<int,int>>로 {{1,1},{0,1}}
        {
            init(1, 1, N, i, { {1,1}, {0,1} });
        }
        else    //b인경우 pair<pair<int,int>,pair<int,int>>로 {{1,0}, {1,1}}해줌
        {
            init(1, 1, N, i, { {1,0}, {1,1} });
        }
    }
    while (Q--)
    {
        int com, L, R; long long A, B;
        cin >> com >> L >> R;
        if (com == 1)
        {
            update(1, 1, N, L, R);
        }
        else
        {
            cin >> A >> B;
            pair<pair<int, int>, pair<int, int>> P = getVal(1, 1, N, L, R);
            long long tempA = (A * P.first.first + B * P.first.second) % mod;
            long long tempB = (A * P.second.first + B * P.second.second) % mod;
            cout << tempA << " " << tempB << '\n';
        }
    }
}

코드가 너무 더럽네..