[Gold IV] 수들의 합 4 - 2015

문제 링크

성능 요약

메모리: 47512 KB, 시간: 204 ms

분류

자료 구조, 해시를 사용한 집합과 맵, 누적 합, 트리를 사용한 집합과 맵

문제 설명

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, k = map(int, input().split())
prefix = list(map(int, input().split()))
arr = [0]*len(prefix)
arr[0] = prefix[0]
for i in range(1, len(prefix)):
    arr[i] = prefix[i] + arr[i-1]

cache = {0:1}
res = 0

for i in range(len(arr)):
    if arr[i]-k in cache.keys():
        res += cache[arr[i]-k]

    if arr[i] not in cache.keys():
        cache[arr[i]] = 1

    else:
        cache[arr[i]] += 1
print(res)

---

회고

맨처음 딕셔너리 생성시 0을 1로 초기화해주지 않아서 중복되는 문제 발생하였음
조금 신경써줄것.
누적합 개념을 처음 공부해봤는데 생각보다 쉬운개념이라 금방 풀이함.
개념만 읽고 점화식을 세운것이 이번 문제 풀이에서 가장 크게 얻어갈 점이었음