[누적합] 2015번 수들의 합 4

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

알고리즘

• 누적합 (구현)

풀이

포인트

• 누적합의 개념을 알고 있어야 하고 코딩할수 있어야 한다.
• key-value 개념을 통해, 이전에 누적합이 i인 것의 개수를 저장한다.

누적합

누적합을 통해 주어진 수열의 부분합을 $O(n)$으로 구할 수 있다. 하지만 이 문제는 부분합이 $k$인 구간이 몇 개 있는지 구해야한다. 수열의 부분합을 구하고 ($O(n)$) $n$개 중 2개를 뽑아서 ($O(n^2))$ 그 부분합이 $k$인지 확인 ($O$(1))하면 된다. 즉, 시간 복잡도는 $O(n^2)$이며, 문제의 $n$의 범위로 보아, 시간 초과가 뜰 것이다.

따라서 이 문제는 시간 복잡도를 $O(n)$ 또는 $O(nlogn)$으로 줄여야 정답 판정을 받을 수 있다. 이를 위해, 파이썬의 key-value 자료 구조를 이용한다.

아이디어는 아래와 같다.

• 0번째 인덱스에서 i번째 인덱스까지의 부분합을 구하며, 그 값이 k가 되면 정답에 더해준다.

• 파이썬의 딕셔너리 자료 구조를 이용하여, dct[i]가 합이 i가 되는, 인덱스 0번째에서 j번째까지의 부분합이 dct[i]개가 되도록 연산한다.

• 0번째 인덱스에서 i번째까지의 부분합을 구할 때마다, dct[ psum[i] - k ]을 정답에 더해준다.

  • 현재 부분합이 모두 0번째 인덱스에서 시작하므로, psum[i] - psum[j] == k를 만족하는 j의 개수를 구하는 것이다.
  • 만약 dct[ psum[i] - k ] = dct[ psum[j] ]가 존재한다면, psum[i] - k값과 동일한 0번째 인덱스 ~ j번째 인덱스까지의 부분합이 존재한다는 뜻이며, 이는 곧바로 j+1 ~ i번째까지의 부분합이 k인 경우의 수와 동일하다.

아래 그림으로 다시 이해해보자.

i번째 까지의 부분합을 구할때마다, dct[ psum[i] - k ]을 정답에 더해준다는 것은, 곧 dct[ psum[j] ]을 더하는 것과 동일하고 결국 j+1 ~ i번째까지의 부분합이 k인 경우를 더하는 것과 동일하다.

[Source Code]

from collections import defaultdict
n,k = map(int,input().split())
a = [0] + list(map(int,input().split()))
s = [0]*(n+1)
ans = 0
dct = defaultdict(int)
for i in range(1,n+1):
    s[i] = s[i-1] + a[i]
    if s[i] == k:
        ans += 1
    ans += dct[s[i] - k]
    dct[s[i]] += 1
print(ans)