[백준] 2015. 수들의 합 4 - Python3

문제 설명

N, K가 주어지고, N개의 수들로 이루어진 배열이 주어질 때, 부분합이 K인 경우의 수를 구하는 문제

https://www.acmicpc.net/problem/2015

2015번: 수들의 합 4

• 알고리즘 분류 - 자료 구조, 누적 합, 해시를 사용한 집합과 맵, 트리를 사용한 집합과 맵
• 난이도 - 골드 4

---

관련 개념

부분합

i~j 번째 요소의 합은 j번째 요소까지의 총합 - i번째 요소까지의 총합과 같다.

---

접근 방법

1. 이중반복문 (실패 - 시간 초과)

배열의 i번째 요소까지의 합을 저장하는 배열 S를 만들고,

0~n-1 범위의 i와 i+1~n-1 범위의 j에 대해 S[i] == k인 경우와 S[j] - S[i] == k인 경우를 카운팅 했다.

 

2. 딕셔너리

(다른 분들의 풀이 참고)

부분합을 한 번 더 응용하여, 새로운 딕셔너리에 각 부분합이 등장하는 횟수를 저장함으로써 1번 풀이에서의 j에 대한 반복문을 작성하지 않고 구할 수 있다.

먼저, j를 i로 대체하기 위해 둘을 바꿔 생각한다. i >= j 일 때, S[i] - S[j] == k를 이항하여, S[j] == S[i] - k 식을 활용함으로써 S[j]를 i와 k에 대한 식으로 대체할 수 있다.

 

1) 0~i번째 요소의 부분합이 k인 경우

1번 풀이와 마찬가지로 단순히 S[i] == k인 경우를 answer에 더한다.

 

2) i~j번째 요소의 부분합이 k인 경우

S[j] = S[i] - k라는 점을 이용하여, 딕셔너리의 키값으로 S[i] - k가 존재하는지 확인한다.

즉, 현재는 i번째 요소를 탐색하는 시점이므로, 이전에 S[i] - S[j] == k를 만족하는 S[j] 값이 있었는지 체크하는 것이다.

만약 있었다면 등장했던 횟수만큼을 answer에 더하고, 없었다면 더하지 않는다(0을 더한다).

---

답안 코드

1. 이중반복문 (시간 초과)

시간 초과

import sys
input = sys.stdin.readline

answer = 0
n, k = map(int, input().split())
A = list(map(int, input().split()))

# S: 0~i번째 요소까지의 부분합 저장
S = [0 for _ in range(n)]
S[0] = A[0]
for i in range(1, n):
    S[i] = S[i-1] + A[i] 

for i in range(n):
	if S[i] == k: # 인덱스 0~i 범위의 합이 k인 경우.
    	answer += 1
    for j in range(i+1,n): # 인덱스 i~j 범위의 합의 k인 경우.
         if S[j]-S[i] == k:
            answer += 1
            
print(answer)

 

2. 딕셔너리

메모리: 47516 KB , 시간: 192 ms, 코드 길이: 524 B

import sys
input = sys.stdin.readline

answer = 0
n, k = map(int, input().split())
A = list(map(int, input().split()))

S = [0 for _ in range(n)]
S[0] = A[0]
for i in range(1, n):
    S[i] = S[i-1] + A[i] 

cnt = {}
for i in range(n):
    # 인덱스 0~i 범위의 합이 k인 경우
    if S[i] == k:
        answer += 1
    
    # 인덱스 i~j 범위의 합이 k인 경우. S[j] = S[i]-k
    answer += cnt.get(S[i]-k,0)
    
    # 인덱스 i까지 탐색했을 시점의 부분합 카운팅
    if S[i] in cnt:
        cnt[S[i]] += 1
    else:
        cnt[S[i]] = 1
            
print(answer)

어떻게 이런 생각을 하지..? 코드 보고 이해하는 것만으로도 한참 걸렸다. 난 똥멍청이다......ㅠㅠ