백준 2015 - 수들의 합 4 (파이썬)

분류 : 누적합

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

풀이

부분합을 활용하여 특정 구간 합이 K가 되는 경우의 수를 찾는 문제이다.

단순한 방법으로는 O(N^2) 시간 복잡도로 모든 구간을 확인 할 수 있지만, 주어진 N의 최대값이 200,000 이므로, 이는 비효율적인 방법이다.

따라서, 누적합과 해시맵을 활용하여 O(N)으로 해결해야한다.

psum = [0] * N
psum[0] = A[0]

for i in range(1, N) :
    psum[i] = psum[i - 1] + A[i]

• psum[i] 는 A[0] 부터 A[i] 까지의 합을 저장하는 누적합 배열이다.
• psum[i] = psum[i - 1] + A[i]를 이용하여 이전 누적합에 현재 값을 다하는 방식으로 누적합을 구한다.
• 이과정은 O(N)에 수행된다.

for i in range(N) :
    goal = psum[i] - K

    if goal == 0:
        answer += 1

    if goal in count :
        answer += count[goal]

    if psum[i] not in count :
        count[psum[i]] = 0
    count[psum[i]] += 1

• goal = psum[i] - K
  • psum[j] - psum[i-1] = K 를 만족하는 j를 찾기 위해 변형하면psum[i-1] = psum[j] - K 가 된다.
  • 즉, goal = psum[i] - K 가 과거의 누적합 중에 존재하는지 찾으면 된다.
• if goal == 0:
  • psum[i] == K 라면 0부터 i까지의 부분합이 K인 경우 이므로 answer를 증가시킨다.
• if goal in count:
  • 과거에 goal = psum[i] - K 인 값이 나온 적이 있다면, 그만큼 answer 를 증가시킨다.
  • 즉, psum[j] = goal 인 j의 개수만큼 정답에 추가된다.
• count[psum[i]] 업데이트
  • 현재 psum[i] 값을 count 에 추가하여, 이후에 같은 값이 나올 경우를 대비한다.

코드

# 백준 2015 - 수들의 합 4
# 분류 : 누적합

N, K = map(int, input().split())
A = list(map(int, input().split()))

psum = [0] * N
psum[0] = A[0]

for i in range(1, N) :
    psum[i] = psum[i - 1] + A[i]

answer = 0
count = {}

for i in range(N) :
    goal = psum[i] - K

    if goal == 0:
        answer += 1

    if goal in count :
        answer += count[goal]

    if psum[i] not in count :
        count[psum[i]] = 0
    count[psum[i]] += 1

print(answer)