[BOJ_2015] 수들의 합 4 (python)

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

 

✅ 문제 설명

수열이 주어졌을 때, 부분합이 K인 것이 몇개나 있는 지 구하는 문제다.

✅ 접근 방식

처음에는 인덱스 0부터 N-1까지 돌면서 해당 인덱스부터 시작하는 배열의 부분합을 구해, K가 나오는 수를 계산하려고 했다.

-> 근데 이 방법은 1 ≤ N ≤ 200,000이므로, O(N^2)이 소요되면 40,000,000,000(4백억..) 절대 2초만에 들어올 수 없을 것이다.

 

그래서 나도 사실 답을 보고 왔다!

알고리즘 분류에 해시, 트리를 사용한 집합과 맵이라 되어있어 이게 뭘까..? 딕셔너리를 사용하는 걸까 싶었는데,

역시 딕셔너리를 사용하더라! 똑똑해..

 

✅ 코드 설명

dict = {0: 1} # 누적합을 기록해둘 딕셔너리

우선 인덱스 0부터 누적합을 계산해나가기 전에, 누적합을 기록해줄 딕셔너리를 선언해둔다.

이때 누적합 값 0이 1개 있다라는 표시를 꼭!! 해줘야 정확한 답이 나온다.

아무것도 안 누적했을 때 0이므로, 이를 꼭 딕셔너리에 값을 넣어주고 다음으로 진행하자!

(이걸 안하면, 2 -2 2 수열이 있다고 쳤을 때 K=2라하면 [2] [2,-2,2] 이렇게 2개의 답이 존재하는데 0:1을 안 넣어주면,, [2]는 답으로 인식하지 못한다.)

 

이후,

sum = 0
count = 0
for i in range(N):
    sum += vals[i]

    # 1. 부분합이 K인 것이 있는 지 확인
    if sum - K in dict: # 지금가지의 누적합에서 K만큼 차이나는게 누적합 기록 딕셔너리에 있으면
        count += dict[sum - K] # 부분합이 K인 것이므로 그 수만큼 더해준다.

    # 2. 지금의 누적합을 딕셔너리에 기록
    if sum in dict:
        dict[sum] += 1
    else:
        dict[sum] = 1

누적합을 계산할 sum 변수와, 부분합이 K인것의 개수를 저장해둘 count 변수를 선언했다.

그리고 인덱스 0부터 N-1까지 돌면서 부분합을 계산한다.

sum += vals[i]

 

1번에서 부분합이 K인것이 있는지 확인해준다. 어떻게?

지금까지의 누적합에서 K만큼 차이가 나는 누적합이 있었는지 확인한다. 즉, sum-K 값을 갖는 누적합이 있었다면

해당 누적합부터 현재까지의 부분합이 K이므로 sum-K값을 가졌던 누적합의 개수만큼 count에 더해준다.

count += dict[sum-K] 를 해준다.

 

2번에서는 지금의 누적합(sum)을 다음을 위해 딕셔너리에 기록해둔다.

sum이 딕셔너리에 있으면 +=1 로 더해주고, 없었다면 1로 새롭게 저장해준다.

 

 

✅ 최종 코드

# 수들의 합 4
# 부분 합 중 합이 K인 것이 몇개나 있는지 구하시오.
import sys
from collections import defaultdict

N, K = map(int, sys.stdin.readline().split())
vals = list(map(int, sys.stdin.readline().split()))

dict = {0: 1} # 누적합을 기록해둘 딕셔너리
# 아무것도 선택하지 않았을 때의 누적합 0은 기본으로 하나 넣어둠.

sum = 0
count = 0
for i in range(N):
    sum += vals[i]

    # 1. 부분합이 K인 것이 있는 지 확인
    if sum - K in dict: # 지금가지의 누적합에서 K만큼 차이나는게 누적합 기록 딕셔너리에 있으면
        count += dict[sum - K] # 부분합이 K인 것이므로 그 수만큼 더해준다.

    # 2. 지금의 누적합을 딕셔너리에 기록
    if sum in dict:
        dict[sum] += 1
    else:
        dict[sum] = 1

print(count)

 

✅ 코멘트

- 해당 문제 알고리즘 분류를 보니 트리르 사용한 집합과 맵이라는 키워드도 있던데, 해시 = 딕셔너리를 이용해서 푸는 방법 말고 트리로 푸는 방법이 있는 지 알아봐야겠다.

- 접근 방식이 생각나지 않아, 답을 참고했다..! 까먹을때 쯔음 다시 풀어봐야겠다.