Leetcode 문제 풀이 (1) - 합이 k 가 되는 부분 배열 (Subarray) 의 갯수

문제 기술 (Problem Statement)

정수 배열 nums 와 어떤 정수값 k 가 주어질 때, 모든 원소 (Element) 들의 합이 k 가 되는 부분 배열 (Subarray) 의 갯수를 구하여라.

문제 유형 (Category)

Prefix Sum

출처 (References)

문제 설명 (Explanation)

다소 낮은 acceptance rate 을 보이는 문제로, brute force (주먹구구식) 형태의 솔루션은 생각해내기 쉬우나, 제일 효율적인 (Optimal) 솔루션을 위해서는 DP (Dynamic Programming) 기법과 수학적 직관이 필요한 문제입니다.

우선 문제 분석부터 해봅시다. 주어진 데이터는 integer array (정수형 배열) 이며, 이 배열에서 각 부분들 (즉, 부분 배열) 의 합이 k 가 되는 이 부분들의 등장 횟수를 구하면 되는 문제죠. 일단 제일 간단한 brute force 를 이용한 솔루션을 살펴봅시다.

Solution 1. Brute Force (Iterative)

단순하게 생각하자면, 모든 부분 배열들의 조합을 다 테스트해서 그 중 합이 k 가 나오는 배열들만 세어주면 되는거죠. 즉, 인덱스를 가리키는 포인터를 i라 할때, i가 0 부터 시작을 해서 배열 끝까지 (n-1, n := nums 의 길이) 순회 (iterate) 하면서 매 iteration 마다 [i, j] 구간 (j := i ... n-1) 의 element 들을 다 합해서 테스트 해주면 되는거죠.

아래는 Java 로 짜여진 코드입니다:

int getNumSubarrays(int[] nums, int k) {
    var count = 0;
    for (var i = 0; i < nums.length; i++) {
        var sum = 0;
    
        for (var j = i; j < nums.length; i++) {
            sum += nums[j];
            
            if (sum == k) {
                count++;
            }
        }
    }
    
    return count;
}

시간 복잡도 (Time Complexity)

(여기서 시행 횟수는 합 (sum) 을 구하는 시행 횟수이자 비교하는 횟수도 포함합니다. 즉, 원칙적으로는 둘다 카운팅 해야하지만 둘다 결국 $ Theta(1) $ 시간 복잡도를 가지니 하나로 쳐서 세어 주겠습니다.)

Iteration #1: i = 0, j = 0 ... n-1 => 시행 횟수: n

Iteration #2: i = 1, j = 1 ... n-1 => 시행 횟수: n -1

...

Iteration #n: i = n-1, j = n-1 ... n-1 => 시행 횟수: 1

$$ T(n) = \sum_{k=n}^{1}k = \sum_{k=1}^{n}k = \frac{n(n+1)}{2} = \theta(n^{2}) $$

(이 때, 항상 $ \frac{n(n+1)}{2} $ 만큼 시행되므로, Theta 표기법을 사용해줄 수 있습니다.)

공간 복잡도 (Space Complexity)

따로 auxiliary space 를 사용하지 않으므로 (sum => 1, count => 1):

$ \theta(1) $

최적화 (Optimization)

위 방법으로는 좀 뭔가 아쉽지 않나요? 우선 패턴을 보자면, 매 iteration 마다 새로 부분합을 계산하고 있죠. 즉, 각 부분 배열의 합을 계산하는 동작이 반복이 된다는 뜻입니다. 즉, 구간 $ [i, j] $ 의 합을 $ s_{(i, j)} = \sum_{k=i}^{j}nums[k] $ 이라 정의할 때, 각 iteration 에서 계산하는 부분 합들을 나열해보면:

Iteration #1: $ s_{(0, 0)}, \; s_{(0, 1)}, \; s_{(0, 2)}, \; \cdots , \; s_{(0, n-1)} $

Iteration #2: $ s_{(1, 1)}, \; s_{(1, 2)}, \; s_{(1, 3)}, \; \cdots , \; s_{(1, n-1)} $

Iteration #2: $ s_{(2, 2)}, \; s_{(2, 3)}, \; s_{(2, 4)}, \; \cdots , \; s_{(2, n-1)} $

...

Iteration #n: $ s_{(n-1, n-1)} $

이고, 이 때 $ s_{(i, j)} = s_{(0, j)} - s_{(0, i)} \; (j \geq i) $ 를 통해 구할 수 있으므로, 결국 $ s_{(0, 0)}, \; s_{(0, 1)}, \; s_{(0, 2)}, \; \cdots , \; s_{(0, n-1)} $ 만 구한다면 나머지 부분 배열의 합들은 전부 이 $ n $ 가지의 합을 서로 조합해서 빼주는 식으로 구할 수 있습니다. 따라서, 이 합들을 저장하기 위해 $ \theta(n) $ 만큼의 공간을 더 사용하겠지만, 매 iteration 마다 합을 새로 구해주는 작업은 안해도 되므로 좀 더 효율적인 알고리즘을 만들 수 있게 되는거죠. 이 때 첫 인덱스부터 i 번째 까지의 누적 합, 즉 $ s_{(0, i)} $ 을 prefix sum 이라고 불러줍니다. (전체 배열에서 i 번째 까지의 prefix 배열에 대한 합을 구해주기 때문에)

따라서, $ \theta(n) $ 만큼의 공간을 더 사용해 각 iteration 에서의 prefix sum 을 새로운 배열에 저장을 함으로써 부분 배열들의 합계를 구하는 작업만큼은 최적화를 해줄 수 있습니다.

prefix sum 을 통해 최적화를 해주더라도 결국 모든 $ (i, j) (j \geq i) $ 쌍의 부분합을 다시 구해야 하므로 최종적인 시간 복잡도는 변하지 않습니다.

Solution 2. Optimization via Prefix Sum (Prefix Sum 을 통한 최적화)

Figure 1. 각 인덱스에서의 Prefix Sum 및 부분배열의 갯수

그렇다면 $ \theta(n^{2}) $ 의 시간 복잡도로 만족을 해야하는걸까요? 조금 더 최적화를 할 수 있지 않을까요? 여기서 부터가 약간의 수학적 트릭이 필요한 부분입니다.

임의의 인덱스 $ i, j $ 사이 (즉 구간 $ (i, j] $) 의 부분 배열의 합을 두 인덱스에 관한 함수 $ S(i, j) = s_{(0, j)} - s_{(0, i)} \; (j \geq i) $ 라고 정의할 때, 위 그림에서 $ (i_{2}, i_{3}] $ 구간의 합, 즉 $ S(i_{2}, i_{3}) = s_{(0, i_{3})} - s_{(0, i_{2})} $ 이 되므로 $ (s + 2k) - (s + k) = k $ 를 통해 이 구간의 합은 $ k $ 가 됨을 알 수 있습니다. 즉 각 index 까지의 prefix sum 을 저장하는 것으로 임의의 부분 배열 $ (i, j] $ 의 합을 구할 수 있게 되는거죠. 또, $ (i_{1}, i_{2}] $ 구간에서는 마찬가지 방식으로 이 구간의 합이 0이 됨을 알 수 있습니다.

마찬가지로 위 방식으로 $ (i_{1}, i_{3}] $ 구간의 부분 배열을 생각해 본다면 $ S(i_{1}, i_{2}) = 0 $ 이었기 때문에 $ S(i_{1}, i_{3}) = S(i_{1}, i_{2}) + S(i_{2}, i_{3}) = 0 + k = k $, 즉 이 부분 배열의 합은 마찬가지로 $ k $ 가 되는 거죠.

이를 일반화 해보면, 임의의 두 인덱스 사이의 구간 (i, j] 의 부분 배열의 합이 $ k $ 가 될때, $ (i, j] $ 내의 합이 0이 되는 부분 배열 갯수만큼 합이 $ k $ 가 되는 부분 배열의 총 갯수가 더해지는 것을 알 수 있습니다.

즉, 구간 $ (i, j] $ 내에서 합이 $ x $ 가 되는 부분 배열의 갯수를 x 에 관한 함수 $ C_{(i, j)}(x) $ 로 정의할때 다음이 성립합니다:

$$ S(i, j) = k \; \rightarrow \; C_{(i, j)}(k) = 1 + C_{(i, j)}(0) $$

이 때, 문제에서 요구하는 값은 전체 배열 ([0:n]) 즉, 구간 $ [0, n) = (-1, 0] + (0, n) = [-1, n-1] $ 에서 합이 $ k $ 가 되는 부분 배열의 갯수이므로, 이를 위의 식을 이용해 적으면 $ C_{(-1, n-1)}(k) = \sum_{(i, j)}^{}C_{(i, j)}(k) \; (-1 \leq i \leq j \leq n-1) $

Learning Never Stops