C++로 풀어보는 이진 배열 삼등분 알고리즘 문제

0과 1로만 이루어진 배열 A가 주어졌다고 가정해 봅시다. 이 배열을 세 개의 비어 있지 않은 부분으로 나누되, 세 부분이 모두 동일한 이진(binary) 값을 나타내도록 만들어야 합니다. 가능하다면 i + 1 < j 조건을 만족하는 인덱스 쌍 [i, j] 중 아무거나 하나를 반환하면 됩니다.

• 첫 번째 부분: A[0], A[1], ..., A[i]
• 두 번째 부분: A[i+1], A[i+2], ..., A[j-1]
• 세 번째 부분: A[j], A[j+1], ..., A[A.length - 1]

세 부분을 같은 이진 값으로 나누는 것이 불가능하다면 [-1, -1]을 반환합니다.

예를 들어 입력이 [0,1,0,1,1]이라면 출력은 [0,4]가 됩니다. 참고로 이 문제는 유효한 답이 여러 개일 수 있으며, 그중 어떤 것을 반환해도 정답으로 인정됩니다.

풀이 접근 방법

이 문제는 보조 함수 getIdx()를 활용해 다음 단계에 따라 해결할 수 있습니다.

getIdx() 함수

1. 배열 a와 인덱스 left, right를 매개변수로 받습니다.
2. left < right이면서 a[left]가 0인 동안 left를 1씩 증가시켜 선행 0을 건너뜁니다.
3. right가 배열의 크기보다 작은 동안 다음을 반복합니다.
   • a[left]와 a[right]가 다르면 -1을 반환합니다.
   • 같다면 left와 right를 각각 1씩 증가시킵니다.
4. 반복이 종료되면 left - 1을 반환합니다.

메인 로직

1. 크기가 2인 배열 ret을 선언하고 -1로 초기화합니다.
2. num := 0, n := 배열 A의 크기로 설정합니다.
3. i를 0부터 n-1까지 순회하며 A[i]가 1일 때마다 num을 1씩 증가시켜 배열 전체에서 1의 개수를 셉니다.
4. num을 3으로 나눈 나머지가 0이 아니라면 세 부분으로 나눌 수 없으므로 ret([-1, -1])을 그대로 반환합니다.
5. num이 0이라면 모든 원소가 0이라는 뜻이므로 {0, 2}를 반환합니다.
6. req := num / 3으로 설정해 각 부분이 포함해야 할 1의 개수를 구합니다.
7. idx := n - 1부터 시작하여 뒤에서부터 1을 req개 셀 때까지 idx를 감소시킵니다. 이를 통해 세 번째 부분의 시작 위치를 파악할 수 있습니다.
8. idx를 1 증가시킨 후, firstEnd := getIdx(A, 0, idx)로 첫 번째 부분의 끝 인덱스를 구합니다. firstEnd가 음수라면 ret을 반환합니다.
9. secondEnd := getIdx(A, firstEnd + 1, idx)로 두 번째 부분의 끝 인덱스를 구합니다. secondEnd가 음수라면 ret을 반환합니다.
10. {firstEnd, secondEnd + 1}을 최종 결과로 반환합니다.

아래 예제 코드를 통해 더 자세히 이해해 보겠습니다.

예제 코드

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
class Solution {
    public:
    vector<int> threeEqualParts(vector<int>& A){
        vector<int> ret(2, -1);
        int num = 0;
        int n = A.size();
        for (int i = 0; i < n; i++) {
            num += (A[i] == 1);
        }
        if (num % 3 != 0)
            return ret;
        if (num == 0) {
            return { 0, 2 };
        }
        int req = num / 3;
        int idx = n - 1;
        for (int temp = 0; idx >= 0 && temp < req; idx--) {
            temp += A[idx] == 1;
        }
        idx++;
        int firstEnd = getIdx(A, 0, idx);
        if (firstEnd < 0)
            return ret;
        int secondEnd = getIdx(A, firstEnd + 1, idx);
        if (secondEnd < 0)
            return ret;
        return { firstEnd, secondEnd + 1 };
    }
    int getIdx(vector<int>& a, int left, int right){
        while (left < right && a[left] == 0)
        left++;
        while (right < (int)a.size()) {
            if (a[left] != a[right])
                return -1;
            left++;
            right++;
        }
        return left - 1;
    }
};
main(){
    Solution ob;
    vector<int> v = {0,1,0,1,1};
    print_vector(ob.threeEqualParts(v));
}

입력

{0,1,0,1,1}

출력

[1, 4]