[백준] 25332. 수들의 합 8 - Python3

문제 설명

길이가 N으로 같은 정수 수열 A, B가 주어질 때, 각 수열의 i~j번째 수의 합이 서로 같은 i,j 쌍의 개수를 구하는 문제

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

25332번: 수들의 합 8

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

---

관련 개념

부분합

---

접근 방법

수들의 합 4 문제와 같은 원리로 풀었다.

 

먼저, A와 B의 부분합이 같은 경우는 아래 두 가지다.

1)  0~i 범위의 합이 같은 경우

이 경우는 단순하다.

A와 B에 대해 부분합 배열을 만들어 A[i] == B[i]인 경우를 카운팅 한다.

 

2) i~j 범위의 합이 같은 경우

A, B 수열에 대해 i~j 범위의 부분합을 a_presum[i:j], b_presum[i:j]라고 할 때, A와 B의 부분합이 같다면 아래와 같은 수식으로 표현할 수 있다.

 

a_presum[i:j] = b_presum[i:j]

a_presum[j] - a_presum[i] = b_presum[j] - b_presum[i]

a_presum[j] - b_presum[j] = a_presum[i] - b_presum[i]

 

i, j에 대한 반복문은 시간복잡도가 O(n^2)으로 시간초과가 나므로, 밑줄 친 j에 대한 식의 결과를 하나의 변수로 나타내고, 이 변수를 딕셔너리에서 찾으면 그 등장횟수가 나온다.

 

즉, A, B 두 수열의 i번째 요소부터 합이 같아지는 j가 몇 개 있는지 알 수 있는 것이다.

 

위 과정을 0~n 범위의 각 i에 대해 수행하면 O(n)의 시간복잡도로 해결할 수 있다.

---

답안 코드

메모리: 86184 KB , 시간: 360 ms, 코드 길이: 620 B

import sys
input = sys.stdin.readline

n = int(input())
A = list(map(int, input().split()))
B = list(map(int, input().split()))

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

answer = 0
cnt = {}
for i in range(n):
    if A_presum[i] == B_presum[i]:
        answer += 1
    
    s = A_presum[i] - B_presum[i]
    answer += cnt.get(s, 0)
    
    if s in cnt:
        cnt[s] += 1
    else:
        cnt[s] = 1
        
print(answer)