[Python/파이썬] 백준 2015번 수들의 합 4

2015번: 수들의 합 4

 

문제

A[1], A[2], ..., A[N]의 N개의 정수가 저장되어 있는 배열이 있다. 이 배열 A의 부분합이란 1 ≤ i ≤ j ≤ N인 정수 i와 j에 대해 A[i]부터 A[j]까지의 합을 말한다.

N과 A[1], A[2], ..., A[N]이 주어졌을 때, 이러한 N×(N+1)/2개의 부분합 중 합이 K인 것이 몇 개나 있는지를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 정수 N과 K가 주어진다. (1 ≤ N ≤ 200,000, |K| ≤ 2,000,000,000) N과 K 사이에는 빈칸이 하나 있다. 둘째 줄에는 배열 A를 이루는 N개의 정수가 빈 칸을 사이에 두고 A[1], A[2], ..., A[N]의 순서로 주어진다. 주어지는 정수의 절댓값은 10,000을 넘지 않는다.

출력

첫째 줄에 합이 K인 부분합의 개수를 출력한다.

코드

처음 시도한 코드 -> 시간 초과

import sys
input = sys.stdin.readline

n, k = map(int, input().split())
arr = list(map(int, input().split()))

ps = [0]
for i in range(1, n+1):
    ps.append(ps[i-1]+arr[i-1])
    
# i ~ j까지의 부분합 : ps[j] - ps[i-1]
cnt = 0
for i in range(1, n+1):
    for j in range(i):
       if ps[i] - ps[j] == k:
           cnt += 1
print(cnt)

예상대로 당연히 시간 초과...위 코드는 시간 복잡도가 O($N^2$)로 N의 범위가 1 ≤ N ≤ 200,000이므로 시간 초과가 발생함은 당연했다. 어떻게 풀어야할지 도저히 모르겠어서 검색해보니 합이 k를 만들 수 있는 횟수만 구하면 되므로 누적합에서 k를 뺀 값을 이전에 만들 수 있었던 횟수를 더하면 되는 것이었다.

from collections import defaultdict
import sys
input = sys.stdin.readline

n, k = map(int, input().split())
arr = list(map(int, input().split()))

ans = 0
ps_dict = defaultdict(int)
tmp = 0
for a in arr:
    tmp += a
    if tmp == k:    # 처음부터 이제까지 누적합이 k라면
        ans += 1
    
    ans += ps_dict[tmp-k]   # 합이 k를 만들 수 있는 구간합이 있다면
    ps_dict[tmp] += 1
print(ans)

tmp에 누적합을 저장하며 진행하는데 만약 tmp 값이 k가 된다면 우선 ans에 1을 더해준다. 그리고 ans에 ps_dict[tmp-k]의 값을 더해주는데 이 값을 왜 더해주냐면 ps_dict[tmp-k]에는 앞의 값들부터 지금 더해진 값(a)까지의 부분합이 k를 만들 수 있는 경우의 수들이 저장되어 있기 때문이다. 언제 만들어지는지는 알 필요 없고 몇 번 만들어지는지만 구하면 되므로 이렇게 풀이할 수 있었다.

이렇게 풀이하니 O(N) 시간만에 풀이할 수 있었다.