빠른 구간 합 구하기(누적 합)

[1차원 리스트의 누적합]

N개의 정수를 원소로 갖는 배열이 주어졌을 때 모든 구간 합을 구하기 위해서는 2중 반복문을 실행하면서 선형탐색 결과를 더해가야 한다.

 

하지만 별도의 누적합계 테이블을 만들어 놓으면 연산을 많이 줄일 수 있다.

 

 

import sys


data = [-40, 20, -10, 10, -30, 10]
table = [0]
for d in data:
    table.append(table[-1] + d)


max = -sys.maxsize
s_idx = 0
e_idx = 0
for i in range(1, len(table)):
    for j in range(i):
        if max < table[i] - table[j - 1]:
            max = table[i] - table[j - 1]
            s_idx = j - 1
            e_idx = i - 1


print(f'최대 부분 합은 {max}이고, 인덱스{s_idx} ~ {e_idx}에 해당하는 값들의 합입니다.')

 

 

 

---

 

 

[2차원 리스트의 누적합]

다음과 같은 6 * 6 크기의 2차원 리스트 L이 있다고 하자.

 

 

L과 동일한 크기의 누적합을 위한 리스트 A를 생성하고 다음과 같이 값을 채워간다.

• 현재 인덱스 (i, j)를 기준으로 (i - 1, j)와 (i, j - 1)에 해당하는 요소가 있는지 확인하여 다음과 같이 누적합 테이블을 작성한다.
  • L[i - 1][j]와 L[i][j - 1]이 모두 존재하지 않는다면 (i, j)까지의 누적합 A[i][j]는 L[i][j] 가 된다.
  • L[i - 1][j]가 존재하고 L[i][j - 1]이 없다면 (i, j)까지의 누적합 A[i][j]는 A[i - 1][j] + L[i][j] 가 된다.
  • L[i][j - 1]이 존재하고 L[i - 1][j]가 없다면 (i, j)까지의 누적합 A[i][j]는 A[i][j - 1] + L[i][j] 가 된다.
  • L[i][j - 1]와 L[i - 1][j]가 모두 존재하면 (i, j)까지의 누적합 A[i][j]는 A[i][j - 1] + A[i - 1][j] + L[i][j] - A[i - 1][j - 1] 이 된다.

예를 들어 L[0][0] ~ L[1][1] 까지의 누적합을 구할 때의 리스트 A값은 아래 그림과 같다.

그림의 두 파란색 요소를 보면

• L[1][0]은 A[0][0] ~ A[1][0] 까지의 합이다.
• L[0][1]은 A[0][0] ~ A[0][1] 까지의 합이다.

따라서 두 요소를 더하면 A[0][0]이 중복으로 더해지게 되므로 모두 더한 값에서 A[0][0]을 빼줘야 한다.

그러므로 A[1][1] = A[0][1] + A[1][0] - A[0][0] + L[1][1] 이 된다.

 

 

L을 바탕으로 작성된 2차원 누적합 리스트 A는 다음과 같다.

 

위의 과정을 python으로 구현해보면 다음과 같다.

def get_acc_sum_table(l):
    inf = int(1e9)
    n = len(l)
    m = len(l[0])
    acc_tbl = [[-inf] * m for _ in range(n)]
    acc_tbl[0][0] = l[0][0]

    for i in range(1, n):
        acc_tbl[i][0] = acc_tbl[i - 1][0] + l[i][0]

    for j in range(1, m):
        acc_tbl[0][j] = acc_tbl[0][j - 1] + l[0][j]

    for i in range(1, n):
        for j in range(1, m):
            acc_tbl[i][j] = acc_tbl[i - 1][j] + acc_tbl[i][j - 1] - acc_tbl[i - 1][j - 1] + l[i][j]

    return acc_tbl


L = [
    [1, 0, 1, 1, 0, 1],
    [0, 1, 1, 0, 1, 1],
    [1, 1, 1, 0, 0, 0],
    [1, 0, 1, 0, 0, 1],
    [0, 1, 1, 0, 1, 1],
    [0, 1, 1, 1, 0, 1]
]

A = get_acc_sum_table(L)

for elem in A:
    print(elem)

 

위의 코드를 실행하면 다음과 같이 출력된다. 그림의 누적합 리스트 A와 동일한 결과다.

[1, 1, 2, 3, 3, 4]
[1, 2, 4, 5, 6, 8]
[2, 4, 7, 8, 9, 11]
[3, 5, 9, 10, 11, 14]
[3, 6, 11, 12, 14, 18]
[3, 7, 13, 15, 17, 22]

누적합 리스트 A[i][j] 값은 L[0][0] ~ L[i][j] 모든 요소 합이라는 의미이다.

 

누적합을 구해놓은 이유는 빠르게 구간합을 구하기 위해서이다.

만약 L[i][j] ~ L[k][l] 까지의 구간합을 한다고 했을 때. i와 j가 모두 0보다 큰 경우라면 다음 순서에 따라 구간합을 구할 수 있다.

 

(i, j) = (2, 2) 이고 (k, l) = (4, 3) 일 때 구간합을 구하는 예를 보도록 하자.

 

L[0][0] ~ L[4][3]의 누적합에서

L[0][0] ~ L[4][1] 까지의 누적합을 빼주고

 

L[0][0] ~ L[1][3] 까지의 누적합을 빼준 후에

 

 

중복해서 빠진 L[0][0] ~ L[1][1]을 한 번 더해주면

 

L[2][2] ~ L[4][3]의 구간합을 구할 수 있다.

 

만약 (i = 0 and j > 0) 이거나 (i > 0 and j = 0) 인 경우는 전체에서 빼야할 구간이 한군데씩 밖에 없다.

위의 내용을 정리해보면 다음과 같다.

# A는 누적합 테이블
if i > 0 and j > 0:
	(i,j) ~ (k, l)의 구간합 = A[k][l] - A[k][j - 1] - A[i - 1][l] + A[i - 1][j - 1]
elif i > 0 and j == 0:
    (i,j) ~ (k, l)의 구간합 = A[k][l] - A[i - 1][l]
elif i == 0 and j > 0:
    (i,j) ~ (k, l)의 구간합 = A[k][l] - A[k][j - 1]
else:
    # 구간의 시작이 (0, 0)이면 (k, l)까지의 전체 합이 구간합이된다.
    (i,j) ~ (k, l)의 구간합 = A[k][l]

 

 

---

 

[2차원 리스트 누적합 - 예제1]

https://www.acmicpc.net/problem/2167

2167번: 2차원 배열의 합

 

import sys


n, m = map(int, sys.stdin.readline().rstrip().split())

tgt = list()
for _ in range(n):
    tgt.append(list(map(int, sys.stdin.readline().rstrip().split())))

inf = -10001
acc_tbl = [[inf] * m for _ in range(n)]
acc_tbl[0][0] = tgt[0][0]

for i in range(1, n):
    acc_tbl[i][0] = acc_tbl[i - 1][0] + tgt[i][0]

for j in range(1, m):
    acc_tbl[0][j] = acc_tbl[0][j - 1] + tgt[0][j]

for i in range(1, n):
    for j in range(1, m):
        acc_tbl[i][j] = acc_tbl[i - 1][j] + acc_tbl[i][j - 1] - acc_tbl[i - 1][j - 1] + tgt[i][j]


k = int(sys.stdin.readline().rstrip())
gugan = list()
for _ in range(k):
    pos = list(map(int, sys.stdin.readline().rstrip().split()))
    gugan.append([p - 1 for p in pos])

for g in gugan:
    from_y, from_x, to_y, to_x = g
    if from_y > 0 and from_x > 0:
        print(acc_tbl[to_y][to_x] - acc_tbl[to_y][from_x - 1] - acc_tbl[from_y - 1][to_x] + acc_tbl[from_y - 1][from_x - 1])
    elif from_y > 0 and from_x == 0:
        print(acc_tbl[to_y][to_x] - acc_tbl[from_y - 1][to_x])
    elif from_y == 0 and from_x > 0:
        print(acc_tbl[to_y][to_x] - acc_tbl[to_y][from_x - 1])
    else:
        print(acc_tbl[to_y][to_x])