문제 소개
0과 1로만 이루어진 배열 arr[]가 주어졌을 때, 각 부분 배열(subarray)이 0만 포함하거나 1만 포함하도록 하는 모든 부분 배열의 개수를 세는 것이 목표입니다. 예를 들어 배열이 [1, 0, 0]이라면, 0으로만 이루어진 부분 배열은 [0], [0], [0, 0]의 세 가지이고, 1로만 이루어진 부분 배열은 [1] 하나입니다.
예제 1
입력 − arr[] = { 0, 0, 1, 1, 1, 0 }
출력 − 0으로만 이루어진 부분 배열: 4개 / 1로만 이루어진 부분 배열: 6개
설명 −
0만 있는 경우: [0], [0], [0], [0, 0] → 총 4개 (arr[0], arr[1], arr[5], arr[0~1])
1만 있는 경우: [1], [1], [1], [1, 1], [1, 1], [1, 1, 1] → 총 6개 (arr[2], arr[3], arr[4], arr[2~3], arr[3~4], arr[2~4])
예제 2
입력 − arr[] = { 1, 0, 1, 0 }
출력 − 0으로만 이루어진 부분 배열: 2개 / 1로만 이루어진 부분 배열: 2개
0만 있는 경우: [0], [0] → 총 2개 (arr[1], arr[3])
1만 있는 경우: [1], [1] → 총 2개 (arr[0], arr[2])
접근 방법
핵심 아이디어는 연속된 동일한 값의 구간을 찾는 것입니다. 길이가 n인 연속 구간에서 만들 수 있는 부분 배열의 개수는 n × (n + 1) / 2라는 등차수열 합 공식으로 계산할 수 있습니다. 예를 들어 연속된 1이 세 개([1, 1, 1]) 있다면, 만들 수 있는 부분 배열은 [1]×3, [1, 1]×2, [1, 1, 1]×1로 총 3 × 4 / 2 = 6개입니다.
배열을 두 번 순회하며 각각 0만 포함하는 부분 배열과 1만 포함하는 부분 배열을 별도로 세고, 연속 구간의 개수를 저장하기 위해 count_0과 count_1 두 카운터를 사용합니다. 알고리즘의 진행 과정은 다음과 같습니다.
숫자 배열 arr[]를 입력받습니다.
함수 sub_zero_one(int arr[], int size)는 배열을 받아 0으로만 이루어진 부분 배열의 개수와 1로만 이루어진 부분 배열의 개수를 계산합니다.
부분 배열 개수를 누적할 변수 temp_0과 temp_1을 초기화합니다.
연속된 0과 1의 개수를 세기 위한 임시 변수 count_0과 count_1을 준비합니다.
for 반복문으로 i = 0부터 size 미만까지 배열을 순회합니다.
첫 번째 순회에서 현재 요소가 1이면 count_1을 증가시킵니다.
요소가 1이 아니라면, 지금까지 쌓인 count_1로 계산한 temp_one_1 = count_1 × (count_1 + 1) / 2를 temp_1에 더하고 count_1을 0으로 초기화합니다.
두 번째 순회에서는 같은 방식으로 0에 대해 count_0, temp_one_0, temp_0 변수를 사용하여 처리합니다.
모든 순회가 끝난 후 마지막 연속 구간이 남아 있다면(카운터가 0이 아니라면) 해당 값을 각각 temp_1과 temp_0에 추가합니다.
두 순회가 모두 끝나면 temp_0과 temp_1에 각각 0만 포함하는 부분 배열과 1만 포함하는 부분 배열의 총 개수가 저장됩니다.
C++ 구현 예제
#include <bits/stdc++.h>
using namespace std;
void sub_zero_one(int arr[], int size){
int count_1 = 0;
int count_0 = 0;
int temp_1 = 0;
int temp_0 = 0;
for (int i = 0; i < size; i++){
if (arr[i] == 1){
count_1++;
}
else{
int temp_one_1 = (count_1) * (count_1 + 1) / 2;
temp_1 = temp_1 + temp_one_1;
count_1 = 0;
}
}
for (int i = 0; i < size; i++){
if (arr[i] == 0)
{ count_0++; }
else{
int temp_one_0 = (count_0) * (count_0 + 1) / 2;
temp_0 = temp_0 + temp_one_0;
count_0 = 0;
}
}
if (count_1){
int temp_one_1 = (count_1) * (count_1 + 1) / 2;
temp_1 = temp_1 + temp_one_1;
}
if (count_0){
int temp_one_0 = (count_0) * (count_0 + 1) / 2;
temp_0 = temp_0 + temp_one_0;
}
cout<<"Subarrays with only 0's are : "<<temp_0;
cout<<"\nSubarrays with only 1's are : "<<temp_1;
}
int main(){
int arr[] = { 0, 0, 0, 1, 1, 0, 1};
int size = sizeof(arr) / sizeof(arr[0]);
sub_zero_one(arr, size);
return 0;
}
실행 결과
위 코드를 실행하면 다음과 같은 출력이 생성됩니다 −
Subarrays with only 0's are : 7
Subarrays with only 1's are : 4
복잡도 분석
배열을 두 번 순회하므로 시간 복잡도는 O(n)이며, 추가로 사용하는 보조 공간은 O(1)로 매우 효율적입니다. 이 접근 방식은 배열의 크기가 커져도 선형 시간 안에 문제를 해결할 수 있습니다.