Truck Driver

 

Problem: https://atcoder.jp/contests/abc430/tasks/abc430_c

C - Truck Driver

 

문제 요약

'a' 와 'b'로만 이루어진 문자열S가 주어진다. (e.g. abbaaabaaba)

문자열 S에서,  'a'가 A개 이상이고, 'b'가 B미만인 부분 문자열의 갯수를 구하여라.

 

예시1)

S	abbaaabaaba
A	4
B	2

정답: 3

부문 문자열1 (4, 8): aaaba

부분 문자열2 (4, 9): aaabaa

부분 문자열3 (5, 9): aabaa

 

예시2)

S	bbbbbbbbbbbbb
A	1
B	2

정답: 0

조건에 맞는 부분 문자열 없음

 

문제 풀이

문제에서 몇 가지 사실을 간파해야한다.

 

① 문자 'a'와 'b'의 갯수는 index가 뒤로 갈수록, 갯수가 늘어난다.

이것은 prefix sum을 사용하여 특정 인덱스에서 'a'나 'b'가 등장한 횟수를 효율적으로 계산할 수 있음을 의미한다.

 

② prefix sum은 정렬되어 있으므로, 이진탐색으로 조건에 맞는 index를 빠르게 찾을 수 있다.

prefix sum에서 구하고 싶은 것(index)은 무엇일까?

prefix sum(a): 'a'가 언제부터 A번 이상 등장하는가? → 해당 인덱스 이후로는 'a'는 항상 A번 이상이다

prefix sum(b): 'b'가 언제까지 B번 미만 등장했는가? → 해당 인덱스까지만 'b'는 B번 미만 등장한 것이다.

 

③  조건에 맞는 부분 문자열을 찾아야 한다. → 탐색 범위를 제한해야 한다.

탐색 범위는 줄어들지만, 찾고자 하는 등장 횟수는 변하지 않는다

 

④ prefix sum에서 탐색 범위를 제한하는 것은, 단순히 범위만 제한한다고 되는 것이 아니다.

제한되는 범위에 맞춰서, 찾고자 하는 값(target)도 변해야 올바른 탐색을 할 수 있다.

 

 

⑤ prefix_sum(a) 와 prefix_sum(b)에서 찾은 인덱스의 공통분모(유효구간)가 문제에서 찾는 답이다.

 

풀이 코드

#include <iostream>
#include <string>
#include <vector>
#include <iomanip>

using namespace std;

int main()
{
    int N, A, B;
    string S;
    cin >> N >> A >> B >> S;

    // prefix sum
    vector<int> ps_a(N + 1, 0);
    vector<int> ps_b(N + 1, 0);
    for (int i = 0; i < N; i++)
    {
        ps_a[i + 1] = ps_a[i] + (S[i] == 'a');
        ps_b[i + 1] = ps_b[i] + (S[i] == 'b');
    }

    long long answer = 0;

    for (int i = 0; i < N; i++)
    {
        // i를 이용해 탐색 범위를 줄여나간다
        // 찾는 target도 탐색 범위에 맞춰 변경해야 한다
        auto it = lower_bound(ps_a.begin() + i, ps_a.end(), ps_a[i] + A);
        int r_a = distance(ps_a.begin(), it);

        it = lower_bound(ps_b.begin() + i, ps_b.end(), ps_b[i] + B);
        int r_b = distance(ps_b.begin(), it) - 1;

        if (r_a <= r_b) // 정답 수 세기
        {
            answer += r_b - r_a + 1;
        }
    }

    cout << answer << "\n";
    return 0;
}