C++로 0, 1, 2의 개수가 같은 부분 문자열 개수 구하기

이 문제에서는 0, 1, 2로만 구성된 문자열 str이 주어집니다. 목표는 0, 1, 2가 각각 같은 개수로 포함된 모든 부분 문자열을 찾아 그 개수를 구하는 것입니다. 예를 들어 str이 "12012"라면 조건을 만족하는 부분 문자열은 "120", "201", "012" 세 가지이므로 답은 3이 됩니다.

예제로 이해하기

입력 − str = "112200120"

출력 − 0, 1, 2의 개수가 같은 부분 문자열의 개수: 5

설명 − 조건을 만족하는 부분 문자열은 다음과 같습니다.

str[0-5]="112200", str[1-6]="122001", str[2-7]="220012", str[0-8]="112200120", str[6-8]="120"

입력 − str = "12012"

출력 − 0, 1, 2의 개수가 같은 부분 문자열의 개수: 3

설명 − 조건을 만족하는 부분 문자열은 다음과 같습니다.

str[0-2]="120", str[1-3]="201", str[2-4]="012"

접근 방식

모든 부분 문자열을 일일이 확인하는 완전 탐색은 O(n³)의 시간이 걸려 비효율적입니다. 대신 접두사 차이(prefix difference) 기법을 활용하면 훨씬 빠르게 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다. 어떤 부분 문자열 안에서 0, 1, 2의 개수가 서로 같으려면, 그 시작 위치 바로 앞 접두사와 끝 위치까지의 접두사에서 (0의 개수 − 1의 개수)와 (0의 개수 − 2의 개수)라는 두 차이값이 동일해야 합니다. 따라서 각 위치마다 이 차이 쌍을 계산하고, 지금까지 등장한 동일한 차이 쌍의 빈도를 더해 주면 됩니다.

정수 문자(0, 1, 2)로 이루어진 문자열을 입력받고 길이를 계산한 뒤 처리 함수에 전달합니다.
조건을 만족하는 부분 문자열의 개수를 저장할 임시 변수 count를 준비합니다.
(0의 개수 − 1의 개수, 0의 개수 − 2의 개수) 쌍을 빈도와 매핑하는 map을 생성하고, 빈 접두사를 나타내는 초기 쌍 (0, 0)에 1을 저장합니다.
FOR 루프를 0부터 문자열 길이까지 반복하면서, str[i]가 '0'이면 0의 개수를, '1'이면 1의 개수를, 그 외('2')라면 2의 개수를 각각 증가시킵니다.
현재 위치에서 zero_one = 0의 개수 − 1의 개수, zero_two = 0의 개수 − 2의 개수를 계산합니다.
두 값으로 쌍을 만들어 map에서 해당 쌍의 빈도를 조회하고, 그 값을 count에 더합니다.
map에서 해당 쌍의 빈도를 1 증가시킨 뒤 다음 문자로 진행합니다.
루프가 끝나면 count를 반환하고 결과를 출력합니다.

예제 코드

#include <bits/stdc++.h>
using namespace std;
int count_string(string str, int length){
   int count = 0;
   map<pair<int, int>, int> map_pair;
   map_pair[{0, 0}] = 1;
   int zero = 0, one = 0, two = 0;
   for (int i = 0; i < length; i++){
      if(str[i] == '0'){
         zero++;
      }
      else if(str[i] == '1'){
         one++;
      }
      else{
         two++;
      }
      int zero_one = zero - one;
      int zero_two = zero - two;
      count += map_pair[{zero_one, zero_two}];
      map_pair[{zero_one, zero_two}]++;
   }
   return count;
}
int main(){
   string str = "112200120";
   int length = str.size();
   cout<<"Count of Substrings with equal number of 0s, 1s and 2s are: "<<count_string(str, length);
   return 0;
}

출력

위 코드를 실행하면 다음과 같은 결과가 출력됩니다 −

Count of Substrings with equal number of 0s, 1s and 2s are: 5

복잡도 분석

문자열을 한 번만 순회하므로 std::map 사용 시 시간 복잡도는 O(n log n)이며, unordered_map을 사용하면 평균 O(n)으로 개선할 수 있습니다. 공간 복잡도는 저장되는 차이 쌍의 수에 비례하여 최대 O(n)입니다.