Algorithm - Prefix Sum은 무엇이고 언제 필요할까요?

algorithm

Prefix sum은 어떤 데이터를 누적하여 합한 것을 말합니다.

예를 들어, Int 배열 arr이 다음과 같이 주어진다면,

int arr[] = new int[] {10, 20, 30, 40, 50, 60};

Prefix sum은 prefixSum[i] = prefixSum[i-1] + arr[i]가 됩니다.

int prefixSum[] = new int[arr.length];
prefixSum[0] = arr[0];
for (int i = 1; i < arr.length; i++) {
    prefixSum[i] = prefixSum[i - 1] + arr[i];
}

prefixSum은 다음과 같이 arr의 누적 합을 보여줍니다.

arr       = [10, 20, 30, 40, 50, 60]
prifixSum = [10, 30, 60, 100, 150, 210]

언제 필요할까?

Prefix sum의 특징을 이용하면 특정 인덱스들의 합을 쉽게 구할 수 있습니다. 알고리즘 Quiz에서 배열의 합을 요구하지만 시간이 오래걸리는 경우 Prefix sum을 이용하는 것을 고려해볼 수 있습니다.

예를 들어, arr의 index 1에서 index 4의 합을 구하는 것은 다음과 같이 for문을 사용하여 구현할 수 있습니다.

int sum = 0;
for (int i = 1; i < 5; i++) {
    sum += arr[i];
}
// sum = 140

위의 구현은 O(n)의 시간이 걸리지만, prefixSum을 이용하면 다음과 같이 O(1)의 시간으로 구현할 수 있습니다.

int sum = prefixSum[4] - arr[0];
// sum = 140

위의 코드를 수식으로 정리하면, index i에서 j의 값의 합은 sum = prefixSum[j] - prefixSum[i-1]이 됩니다.

관련 알고리즘 문제

Prefix sum과 관련된 알고리즘 문제로, Leetcode에 560. Subarray Sum Equals K가 있습니다.

인자 k와 일치하는 sub-arrays의 합을 찾고, 그 개수를 구하는 문제입니다.

Given an array of integers and an integer k,
you need to find the total number of continuous subarrays whose sum equals to k.

Example 1:
Input:nums = [1,1,1], k = 2
Output: 2

Note:
The length of the array is in range [1, 20,000].
The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].

문제 풀이 (Brute force)

560. Subarray Sum Equals K 문제를 먼저 풀어보세요. (아래는 Leetcode에서 제공하는 해답을 옴겼습니다.)

아래 코드처럼 Brute force로 문제를 해결할 수 있지만 O(n3)의 시간이 걸립니다.

public class Solution {
    public int subarraySum(int[] nums, int k) {
        int count = 0;
        for (int start = 0; start < nums.length; start++) {
            for (int end = start + 1; end <= nums.length; end++) {
                int sum = 0;
                for (int i = start; i < end; i++)
                    sum += nums[i];
                if (sum == k)
                    count++;
            }
        }
        return count;
    }
}

문제 풀이 (Prefix sum)

Prefix sum를 이용하면 O(n2)의 시간으로 줄일 수 있습니다. Prefix sum으로 for를 하나 제거할 수 있었습니다.

public class Solution {
    public int subarraySum(int[] nums, int k) {
        int count = 0;
        int[] sum = new int[nums.length + 1];
        sum[0] = 0;
        for (int i = 1; i <= nums.length; i++)
            sum[i] = sum[i - 1] + nums[i - 1];
        for (int start = 0; start < nums.length; start++) {
            for (int end = start + 1; end <= nums.length; end++) {
                if (sum[end] - sum[start] == k)
                    count++;
            }
        }
        return count;
    }
}

또한, Map을 사용하면 O(n)의 시간으로 줄일 수도 있습니다.

고민해보시고, Leetcode에서 풀어보세요.

참고

Loading script...
codechachaCopyright ©2019 codechacha