누적합(백준 알고리즘 2015번)

♣ 누적합

https://jy-deeplearning.tistory.com/175

누적합(백준 알고리즘 10986번)

 

♣ 문제 이해

 

이 문제는 정수 N개로 이루어진 배열의 부분합에서 합이 K가 되는 부분합의 개수를 구하는 것이다.  

n = 6
k = 5
list = [1, 2, 3, 4, 5, 0]

 

위 예시를 보면 주어진 리스트 안의 숫자는 4개이고, 이 리스트의 부분합 중 합이 5인 갯수를 찾는다. 

 

list = [1, 2, 3, 4, 5, 0]
#부분합이 5인 조합
[(2, 3), (5), (5, 0)]

 

부분합 5를 만족하는 경우의 수는 위와 같다. 

 

이 문제를 풀기 위해 순서를 아래와 같이 정리했다. 

• 리스트의 누적합을 리스트로 만든다. 
• 여러 개의 누적합 중 두 개를 골랐을 때 차이가 k가 되는 조합을 찾는다.
• 차이가 k인 두 개의 리스트가 있다면 두 리스트의 차가 k인 것이므로 답에 1을 더한다. 

list = [1, 2, 3, 4, 5, 0]
#누적합 리스트
sum_list = [1, 3, 6, 10, 15]
#첫 번째 누적합 = 1, 1과 k(5)만큼 차이나는 조합 => 6 
1 = [1]
6 = [1, 2, 3]

 

누적합 1과 6의 리스트 차이에 해당하는 [2, 3] 의 부분합이 5이다. 이런 식으로 차이가 k인 두 개의 리스트 조합을 다 구하면 부분합이 k인 경우의 수를 모두 구할 수 있는 것이다. 

 

하지만 주어진 숫자의 수가 많아질수록 선택할 수 있는 두 개의 누적합 경우의 수가 너무 많아진다. 따라서 for문을 여러 번 사용하지 않고 문제를 해결할 수 있는 방법을 고민했다. 그리고 누적합의 수를 딕셔너리의 key로 이용함으로써 for문을 돌지 않고도 k만큼 차이나는 경우의 수를 찾을 수 있도록 했다. 

 

import sys
input = lambda: sys.stdin.readline().rstrip()

n, k = map(int, input().split())
numbers = list(map(int, input().split()))
answer = 0
total = 0
dict = {0:1}

 

answer은 부분합이 k인 경우의 수이고, total은 numbers를 하나씩 가져와서 더할 누적합이다. dict는 새로운 누적합이 생길 때마다 이를 key로 추가할 딕셔너리이다. 여기서 이미 dict = {0:1} 이라고 정의했는데 그 이유는 뒤에서 설명하겠다. 

 

#numbers 안의 인자들을 하나씩 가져와서 더해가며 새로운 누적합을 구한다. 
for number in numbers:
    total += number 
    
    if total - k in dict.keys():
        answer += dict[total-k]
        
    if total in dict.keys():
        dict[total] += 1 
    else:
        dict[total] = 1

 

위의 코드를 보면 numbers에서 숫자를 하나씩 가져와서 누적합을 구하고 있다. 그런 뒤, 해당 누적합과 k만큼 차이나는 누적합이 이미 dict에 존재하는지 확인한다.dict에 존재한다면 그 값만큼을 answer에 더한다. 

그리고 새롭게 구한 total을 dict의 key로 추가한다. 이미 total이 key에 포함되어 있다면 value에 1을 더하고 존재하지 않았다면 value 값을 1로 정한다. 

 

이때 내가 궁금했던 점이 두 가지 있다.

첫 번째는 왜 처음에 dict = {0:1} 이라고 정의했는지이다. 아래 예시를 보면 이해하기 쉽다.

 

n = 6, k = 6 
numbers = [6, -6, 0, -6, 0, 0]

# 첫 번째 인자를 누적했을 때 
total = 6

 

위의 numbers의 첫 번째 인자를 누적하면 total=6이다. 이는 k=6이라는 조건에 맞기 때문에 answer에 1을 더해야 한다.

if total-k in dict.keys():
    answer+=dict[total-k]
    
# total = 6, k=6 이므로 total-k = 0

 

위의 if문에 의하면, total = 6이고 k=6이기 때문에 total-k = 0 이다. 맨 처음에 dict = {0:1} 로 정의한 덕분에 answer += dict[0] 에서 answer += 1 이 되고, 그 결과 answer은 1이다. 

 

즉 해당 total(누적합) 스스로가 k인 경우, 이 누적합에서 k를 뺐을 때 0이 되고, 그러면 answer 에 1을 추가해야 하기 때문에 맨 처음에 dict = {0:1} 로 설정한 것이다. 그렇게 해서 스스로가 k인 부분합을 answer에 추가한다. 

 

두 번째 궁금한 점은 현재 누적합과 k만큼 차이나는 누적합 모두를 고려하는 것이라면 왜 total-k 만 고려하고 total+k는 고려하지 않았냐는 것이다. 이 역시 위의 예시를 사용하면 쉽게 알 수 있다. 

 

n=6 k=6
list = [6, -6, 0, 6, 0, 0]
dict = {0:1}

# 첫 번째 인자를 고려한 누적합
total = 6
total - k = 0
answer += dict[0] = 1
dict = {0:1, 6:1}

# 두 번째 인자를 고려한 누적합
total = 0
total - k = -6 
answer += dict[-6] 인데 이는 존재하지 않음
answer = 1 
dict = {0:2, 6:1}

# 여기서 만약 total+k 까지 고려하게 되면
total + k = 6 

# dict[6] = 1 인데, 이는 첫 번째 인자를 고려한 누적합에 해당한다. 
# 즉 [6] 과 [6, -6, 0] 사이의 차이를 고려하게 되는 것
# 이 차이에 해당하는 값은 [-6, 0] 이므로 누적합이 6이 아니라 -6이다. 
# 따라서 total+k 와 total 사이의 차이 리스트의 누적합은 -k 이므로 고려하면 안된다.

 

 

♣ 전체 코드

import sys
input = lambda: sys.stdin.readline().rstrip()

n, k = map(int, input().split())
numbers = list(map(int, input().split()))
answer = 0
total = 0
total_dict = {0:1}

for i in numbers:
	total += i 
		
	if total - k in total_dict.keys():
		answer+= total_dict[total-k]
		
	if total in total_dict.keys():
		total_dict[total]+=1 
	else:
		total_dict[total] = 1

print(answer)