백준 2015번 파이썬 문제풀이(수들의 합 4)

---

해당문제는 누적합의 응용 문제라고 생각하면 된다. 

 

우선 누적합을 구해놓고 누적합을 처음부터 N 까지 돌면서 해당 값의 찾고하는 값 K 를 뺀 cnt리스트의 인덱스에 저장되어 있는 값을 우선 더해준다.(즉, cnt[prefix[i] - K]) 그 후 해당 누적합의 개수를 cnt 리스트의 해당 인덱스에 +1개 추가해준다. (즉, cnt[prefix[i]] += 1)

 

하지만 K의 값이 아주크거나 음수가 존재할 경우 cnt의 배열크기가 아주 커지게 된다. 혹은 음수 계산을 할 수가 없다. 따라서 딕셔너리 형태로 변환하여 위의 상황을 진행해야 한다.

 

N, K = map(int, input().split())
arr = [0] + list(map(int, input().split()))  # 입력받은 수들을 저장

prefix = [0]  # 누적합을 저장할 리스트

for i in range(1, N + 1):
    prefix.append(prefix[i - 1] + arr[i])

# 만약 모든 값들이 양수이고(누적합이 계속해서 증가) K의 값이 적당할 때는 리스트를 만들어 다음과같이 사용
# cnt = [0 for _ in range(100010)]  # 값의 개수를 저장할 리스트
# ans = 0
# for i in range(1, N + 1):
#     ans += cnt[prefix[i] - K]
#
#     cnt[prefix[i]] += 1
#
# print(ans)

# 만약 음수가 존재하고 K의 값이 아주 큰값일 때는 딕셔너리를 만들어 사용
cnt = {}
ans = 0
for i in range(N + 1):
    ans += cnt.get(prefix[i] - K, 0)

    cnt[prefix[i]] = cnt.get(prefix[i], 0) + 1

print(ans)