0과 1로만 이루어진 배열 A가 주어졌을 때, 원소의 합이 정확히 S가 되는 비어 있지 않은 부분 배열의 개수를 구하는 문제입니다.
예를 들어 입력이 [1,0,1,0,1]이고 S = 2라면, 답은 4가 됩니다. 조건을 만족하는 부분 배열은 다음과 같습니다.
- [1, 0, 1] (인덱스 0~2)
- [1, 0, 1, 0] (인덱스 0~3)
- [0, 1, 0, 1] (인덱스 1~4)
- [1, 0, 1] (인덱스 2~4)
접근 방법: 슬라이딩 윈도우
이 문제는 슬라이딩 윈도우 기법을 응용하면 O(N) 시간에 해결할 수 있습니다. 핵심 아이디어는 '합이 x 이하인 부분 배열의 개수'를 세는 atMost(x) 함수를 만들고, 이를 활용해 정확히 S가 되는 경우의 수를 계산하는 것입니다.
구체적인 단계는 다음과 같습니다.
- 배열 A와 정수 x를 인자로 받는
atMost()메서드를 정의합니다. - x < 0이면 즉시 0을 반환하고, 그렇지 않으면 j := 0, ret := 0으로 초기화합니다.
- i를 0부터 배열 A의 크기까지 순회합니다.
- x에서 A[i]를 뺍니다.
- x < 0이 되는 동안 x에 A[j]를 더하고 j를 1씩 증가시켜 윈도우의 왼쪽 끝을 오른쪽으로 이동시킵니다.
- 현재 i를 오른쪽 끝으로 하는 유효한 부분 배열의 개수는 (i − j + 1)이므로 ret에 더해줍니다.
- 순회가 끝나면 ret을 반환합니다.
- 메인 메서드에서는 다음과 같이 계산합니다.
- ret := atMost(A, S) − atMost(A, S − 1)
- ret을 반환합니다.
'합이 S 이하인 부분 배열 개수'에서 '합이 S − 1 이하인 부분 배열 개수'를 빼면, 자연스럽게 '합이 정확히 S인 부분 배열 개수'만 남게 됩니다.
C++ 구현 예제
#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
int atMost(vector<int>& A, int x){
if(x < 0) return 0;
int j = 0;
int ret = 0;
for(int i = 0; i < A.size(); i++){
x -= A[i];
while(x < 0){
x += A[j];
j++;
}
ret += i - j + 1;
}
return ret;
}
int numSubarraysWithSum(vector<int>& A, int S) {
return atMost(A, S) - atMost(A, S - 1);
}
};
main(){
vector<int> v1 = {1,0,1,0,1};
Solution ob;
cout << (ob.numSubarraysWithSum(v1, 2));
}입력
[1,0,1,0,1]
출력
4
복잡도 분석
atMost() 함수 내부에서 포인터 i와 j는 각각 배열을 한 번씩만 순회하므로, 전체 시간 복잡도는 O(N), 추가 공간 복잡도는 O(1)입니다. 누적 합(prefix sum)과 해시 맵을 사용하는 방법(O(N) 시간, O(N) 공간)에 비해 메모리를 절약할 수 있다는 장점이 있습니다.