[Algorithm] 배열의 구간 합 알고리즘 (Prefix Sum) - Java

구간 합

구간 합(Prefix Sum)은 주어진 배열의 특정 구간, 즉 배열의 연속된 요소들의 합을 의미하다. 예를 들어, 배열이 [1, 2, 3, 4, 5]라 할 때 2번째 요소부터 4번째 요수까지의 구간 합은 2 + 3 + 4로 9가 된다.

 

자바에서 구간 합을 구하기 위해 일반적으로 다음과 같이 반복문을 사용한다.

int[] arr = {1, 2, 3, 4, 5};
int sum = 0;
for (int i = 1; i <= 3; i++) {
    sum += arr[i];
}
System.out.println("sum: " + sum); // 9

 

다음과 같이 2차원 배열의 경우에는 2중 반복문을 통해 구한다.

1 2 3
4 5 6
7 8 9

 

int[][] matrix = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
int sum = 0;
for (int i = 0; i <= 1; i++) {
    for (int j = 0; j <= 1; j++) {
        sum += matrix[i][j];
    }
}
System.out.println("sum: " + sum); // 12

 

 

구간 합 알고리즘

구간 합 알고리즘은 주어진 배열의 특정 구간의 합을 빠르게 계산하는 알고리즘이다. 전체 배열의 누적 합을 미리 계산하여 배열을 만든 뒤, 임의의 구간 합을 O(1)의 시간 복잡도로 빠르게 계산하는 데 사용한다.

 

누적합과 구간합을 구하는 방법은 다음과 같다.

누적합

sum[i] = arr[0] + arr[1] + ... + arr[i]

 

 

누적합의 i번째 요소가 배열의 0번째 요소부터 i번째 합인 것을 활용하여 다음과 같이 쉽게 구할 수 있다.

int[] arr = {1, 2, 3, 4, 5};
int[] sum = new int[arr.length + 1];

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

System.out.println("sum = " + Arrays.toString(sum)); // [0, 1, 3, 6, 10, 15]

 

 

구간합

배열의 i부터 j까지의 구간합은 j번째 누적합에서 i 직전의 누적합을 빼면 된다.

s(i, j) = sum[j] - sum[i - 1]

int[] arr = {1, 2, 3, 4, 5};
int[] sum = new int[arr.length + 1];

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

int i = 2; // 2번쨰 부터
int j = 5; // 5번째 까지의 합
int result = sum[j] - sum[i - 1];

System.out.println("result = " + result); // 14

 

 

2차원 배열의 구간합

2차원 배열에서의 누적합과 구간합을 구하는 방법은 1차원 배열의 구간합 알고리즘을 응용하여 구할 수 있다.

 

다음과 같이 1부터 9까지 3x3 배열의 크기인 2차원 배열이 존재할 때 누적합과 구간합은 다음과 같이 구할 수 있다.

1 2 3
4 5 6
7 8 9

 

 

누적합

2차원 배열의 누적합은 각 위치를 기준으로 (0, 0)부터 현재 위치까지의 모든 요소의 합이다. 조금 복잡할 수 있으나, i, j 번째의 누적합의 공식은 다음과 같다.

sum[i][j] = sum[i][j - 1] + sum[i - 1][j] - sum[i - 1][j - 1] + arr[i - 1][j - 1];

 

• 위 공식에서 sum[i][j]는 배열 arr의 0, 0부터 (i - 1, j - 1)까지의 모든 요소의 합을 나타낸다.
• sum[i][j - 1]은 왼쪽 요소까지의 누적 합이고, sum[i - 1][j]는 바로 위 요소까지의 누적합이다.
• 이를 더하면 왼쪽 위 요소(sum[i - 1][j - 1])까지의 누적합이 2번 포함되므로 sum[i - 1][j - 1]을 한 번 빼준다.
• 이후 현재 위치의 값을 더하면 현재 위치까지의 누적합을 구할 수 있다.

위 공식을 토대로 다음과 같이 코드를 작성할 수 있게 된다.

int[][] arr = {
        {1, 2, 3},
        {4, 5, 6},
        {7, 8, 9},
};
int[][] sum = new int[arr.length + 1][arr[0].length + 1];

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

 

코드의 결과는 다음과 같다.

0  0  0  0
0  1  3  6
0  5 12 21
0 12 27 45

 

구간합

위에서 구한 2차원 배열의 누적합을 토대로 구간합을 구하는 공식은 다음과 같다.

result = sum[x2][y2] - sum[x1 - 1][y2] - sum[x2][y1 - 1] + sum[x1 - 1][y1 - 1]

 

1부터 9까지 3x3 크기의 2차원 배열을 기준으로 보면, 2, 2 위치의 누적합인 12(1 + 2 + 4 + 5)를 구한다고 가정해 보자.

• (1, 1)부터 (1, 2)까지의 누적합 3과 (1, 1)부터 (2, 1)까지의 누적합 5를 더하면 8이 된다.
• 이는 (1, 1)이 2번 더해진 결과이기 때문에 (1, 1)을 한 번 뺀다.
• 마지막으로 (2, 2)에 해당하는 5를 더한다.

 

앞에서 구한 누적합을 통해 (1, 1)부터 (3, 3)까지의 구간합을 구한다면 다음과 같이 구할 수 있다.

• (1, 1)부터 (3, 3)까지의 누적합에서 바로 위까지의 누적합과 바로 왼쪽의 누적합을 뺀다.
• 이는 왼쪽 위 위치까지의 누적 합이 2번 빠지게 되므로 한 번 더한다.

int x1 = 1;
int y1 = 1;
int x2 = 3;
int y2 = 3;
int result = sum[x2][y2] - sum[x1 - 1][y2] - sum[x2][y1 - 1] + sum[x1 - 1][y1 - 1]; // 45