[BOJ/python]2015번 수들의 합 4

백준 2015번 수들의 합 4 파이썬

문제 링크 - https://www.acmicpc.net/problem/2015

2015번: 수들의 합 4

연속하는 부분수열의 합이 K가 되는 구간의 개수를 찾는 문제였다. 주어진 수열의 순서를 바꿀 수 없었기 때문에 누적합을 우선 구하고 시작했다.

그리고 만약 수열의 크기가 작았다면 어떻게 접근했을 지 고민했다. 누적합 배열의 원소들 중 2개를 골랐을 때 차가 K인 경우의 수를 완전 탐색으로 구했을 것이다. 완전 탐색의 시간 복잡도를 줄이기 위한 방법을 생각해보면 굳이 모든 인덱스를 O(N**2) 시간 동안 탐색할 필요가 없다는 것을 알 수 있다.

누적합 배열의 i번째 인덱스의 값이 SUM일 때, 누적합의 값이 SUM+K인 인덱스의 개수만 알면 되기 때문이다.

1. 풀이

• 누적합 배열을 저장한다.
• 마지막 인덱스에서부터 누적합 배열을 탐색한다.
• i번째 인덱스의 누적합의 값이 SUM이라면 딕셔너리에 Key값이 SUM+K인 배열의 길이를 answer에 더해준다.
  • 인덱스의 순서가 역방향이므로 딕셔너리에 저장된 인덱스는 현재 인덱스보다 뒤에 존재한다.
• 만약 현재 누적합의 값이 K라면 answer에 1을 더한다.
• 딕셔너리에 누적합의 값을 Key값으로, 인덱스를 value에 배열로 저장한다.

2. 코드

from collections import defaultdict

N, K = map(int, input().split())
L = list(map(int, input().split()))
S = [0]
for i in range(N):
    S.append(S[-1]+L[i])

answer = 0
idx_dict = defaultdict(list)
for i in range(N,0,-1):
    SUM = S[i]
    
    if SUM == K :
        answer+=1
    

    target = SUM+K
    answer += len(idx_dict[target])
    
    idx_dict[SUM].append(i)
    
print(answer)