[알고리즘/ PYTHON] 누적합 알고리즘

waterlyn·2022년 8월 20일
1
post-thumbnail

출처 및 참고

[알고리즘] 부분합, 누적합 (Prefix Sum) 쉽게 알아보기(파이썬)

일반적으로 부분배열의 합을 구하면 O(N)의 시간이 걸리는데, 누적합을 사용하면 O(1)의 시간으로 구할 수 있다!

1차원 배열에서

준비물

  • 더할 숫자가 담긴 배열 arr
  • 누적합을 담을 배열 sum
arr = [2, 4, 1, -5, 2, -3]
sum = [0, 2, 6, 7, 2, 4, 1]

1차원 배열은 매우 간단하다.

0번째 합은 아무것도 안 더한 값이므로 sum[0]은 0이 되고, 1번째 자리부터 누적합이 차례대로 들어가게 된다.

1차원 배열 누적합 활용 ⭐️

i항부터 j항까지 부분배열에 대한 누적합을 S(i, j)라고 하면 S(i, j)는 다음과 같다.

👉🏻 S(i, j) = sum[j+1] - sum[i]

EX) 2번째부터 5번째까지 누적합 → i = 1, j = 4

S(0, 1) = 2 → sum[1]

S(0, 4) = 2 + 4 + 1 + -5 + 2 = 4 → sum[5]

S(1, 4) = sum[5] - sum[1] = 4 - 2 = 2

2차원 배열에서

준비물

  • 더할 숫자가 담긴 2차원 배열 arr
  • 누적합을 담을 배열 sum_arr
arr = [[1, 2, 3, 4],
			 [2, 3, 4, 5],
		   [3, 4, 5, 6],
		   [4, 5, 6, 7]]
sum_arr = [[0, 0, 0, 0, 0],
		   [0, 1, 3, 6, 10],
		   [0, 3, 8, 15, 24],
		   [0, 6, 15, 27, 42],
		   [0, 10, 24, 42, 64]]

결론부터 말하자면 sum_arr[i, j]에는 arr[0][0]부터 arr[i-1][j-1]까지의 합이 담겨있다는 것이다.

(그림은 블로그 참조) https://yiyj1030.tistory.com/489

위의 사실을 바탕으로 sum 배열을 완성하는 코드는 다음과 같다.

arr = [[1, 2, 3, 4],
			 [2, 3, 4, 5],
		   [3, 4, 5, 6],
		   [4, 5, 6, 7]]
m = 4
n = 3

sum_arr = [[0 for _ in range(m+1)] for _ in range(n+1)]

for i in range(1, n+1):
	for j in range(1, m+1):
		sum_arr[i][j] = arr[i-1][j-1] + sum_arr[i-1][j] + sum_arr[i][j-1]
											- sum_arr[i-1][j-1]
  • sum_arr[i-1][j] + sum_arr[i][j-1]sum_arr[i][j]까지의 합인데 arr[i][j]의 값이 없고, sum_arr[i-1][j-1]까지의 합이 두번씩 합해진 값이 된다.
  • 그래서 식에서 sum_arr[i-1][j-1] 값을 한 번 빼주고, arr[i][j] 값을 더해주는 것.

2차원 배열 누적합 활용 ⭐️

1차원 배열과 동일하게 부분 구간에 대한 누적합을 구하는 것도 가능하다.

arr의 (x1, y1) 부터 (x2, y2)까지의 누적합을 S라고 한다면 다음과 같이 수식이 성립한다.

S = sum_arr[x2+1][y2+1] - sum_arr[x1][y2+1] - sum_arr[x2+1][y1] + sum_arr[x1][y1]
  • (x2, y2)까지의 누적합에서 제외되는 부분의 누적합을 빼는데 교차되는 부분을 두번 빼게 되므로 한번 더해준다! (그림은 블로그 참조) https://yiyj1030.tistory.com/489
  • 제외되는 부분 두개 → sum_arr[x1][y2+1], sum_arr[x2+1][y1]
  • 교차되는 부분 → sum_arr[x1][y1]

EX) (x1, y1) = (1, 1), (x2, y2) = (3, 2)

  • sum_arr[4][3] = arr[0][0]부터 arr[3, 2]까지의 합 = 1 + 2 + 3 + 2 + 3 + 4 + 3 + 4 + 5 + 4 + 5 + 6 = 42
  • sum_arr[1][3] = arr[0][0]부터 arr[0][2]까지이 합 = 1 + 2 + 3 = 6
  • sum_arr[4][1] = arr[0][0]부터 arr[3][0]까지의 합 = 1 + 2 + 3 + 4 = 10
  • sum_arr[1][1] = arr[0][0] = 1
profile
Hello there 🖤

0개의 댓글