문제 이해하기
0과 1로만 구성된 배열 arr[]가 주어졌을 때, 0과 1의 개수가 서로 같은 모든 부분 배열(subarray)의 개수를 세는 것이 목표입니다. 예를 들어 배열이 [1,0,0]이라면 조건을 만족하는 부분 배열은 [1,0] 하나뿐입니다.
예제
입력 − arr[] = { 0, 0, 1, 1, 1, 0 };
출력 − 1과 0의 개수가 같은 부분 배열의 개수: 4
설명 − 조건을 만족하는 부분 배열은 다음과 같습니다.
arr[0~3] = [0,0,1,1],
arr[1~2] = [0,1],
arr[4~5] = [1,0],
arr[0~5] = [0,0,1,1,1,0]
입력 − arr[] = { 0, 1, 1, 1, 1 };
출력 − 1과 0의 개수가 같은 부분 배열의 개수: 1
설명 − 조건을 만족하는 부분 배열은 arr[0~1] = [0,1] 하나입니다.
접근 방법
두 개의 for 루프를 사용해 가능한 모든 부분 배열을 생성하며 배열을 순회합니다. 바깥 루프의 인덱스 i는 0부터 size-1까지, 안쪽 루프의 인덱스 j는 i부터 size-1까지 반복하여 arr[i]부터 arr[j]까지의 부분 배열을 만듭니다. 각 부분 배열에서 0과 1의 등장 횟수를 각각 계산하고, 두 값이 같으면 카운트를 증가시킵니다.
숫자로 이루어진 배열 arr[]를 준비합니다.
함수 sub_zeroes_ones(int arr[], int size)는 배열을 받아 0과 1의 개수가 같은 부분 배열의 개수를 반환합니다.
초기 카운트 값을 0으로 설정합니다.
i=0부터 i<=size-1까지, j=i부터 j<=size-1까지 두 개의 for 루프로 배열을 순회합니다.
부분 배열 arr[i]부터 arr[j]까지의 0과 1의 개수를 저장할 변수 total_0, total_1을 0으로 초기화합니다.
arr[j]의 값을 확인합니다. arr[j]가 0이면 total_0을, 1이면 total_1을 증가시킵니다.
total_0 == total_1이면 카운트를 증가시킵니다. 이는 해당 부분 배열에 포함된 0과 1의 개수가 같다는 의미입니다.
모든 루프가 종료되면 count를 결과로 반환합니다.
예제 코드
#include <bits/stdc++.h>
using namespace std;
int sub_zeroes_ones(int arr[], int size){
int count = 0;
for (int i = 0; i <= size - 1; i++){
int total_0 = 0;
int total_1 = 0;
for (int j = i; j <= size - 1; j++){
if (arr[j] == 0){
total_0++;
}
else if (arr[j] == 1){
total_1++;
}
if(total_0 == total_1){
count++;
}
}
}
return count;
}
int main(){
int arr[] = {0, 1, 1, 0, 0};
int size = sizeof(arr)/sizeof(arr[0]);
cout<<"Count of subarrays with equal number of 1's and 0's are: "<<sub_zeroes_ones(arr, size);
}
실행 결과
위 코드를 실행하면 다음과 같은 출력이 생성됩니다.
Count of subarrays with equal number of 1's and 0's are: 4
시간 복잡도 참고
위 방법은 두 개의 중첩 루프를 사용하므로 시간 복잡도는 O(n²)입니다. 더 큰 입력에 대해서는 0을 -1로 치환한 뒤 누적 합(prefix sum)과 해시 맵을 활용하면 O(n) 시간에 문제를 해결할 수 있습니다.