[Python] 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인 부분합의 개수를 출력한다.

해설

누적합이 진짜 뉴비 분쇄기인듯...

일단 N^2 시간복잡도는 바로 시간초과(10^10)라서 제외했고 누적합을 사용해서 풀어봤는데.. 틀렸습니다가 떠서 다른 방법으로 질문게시판 뒤져보면서 해답을 찾았다.

0 <= i <= j < n 일때,
결국 구간 i - 구간 j = k 인 값을 찾는 것이니,
구간 i - k = 구간 j 를 찾아서 세어주면 답이 된다.

중복 값에 대한 구분 없이 set으로 확인하니 틀렸습니다.가 뜬 것이었다.
매번 try except 쓰는 것도 귀찮아서 defaultdict 알아보고 사용했다.

구간합 문제는 코딩 문제라기 보다는 수학적인 접근을 더 우선적으로 생각해보는게 좋다는 걸 생각했다.

소스 코드

from collections import defaultdict

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

count = 0
prefix_sum = 0
freq = defaultdict(int)
freq[0] = 1  # 누적합 0이 한 번 등장(빈 구간)

for num in q:
    prefix_sum += num
    count += freq[prefix_sum - k]  # (현재 누적합 - k)가 이전에 등장했으면 카운트 증가
    freq[prefix_sum] += 1  # 현재 누적합 등장 횟수 갱신

print(count)