[백준 2143 파이썬] 두 배열의 합 (골드 3, 누적 합)

배코딩·2025년 2월 7일
0

PS(백준)

목록 보기
123/131

소스 코드

import sys
from collections import defaultdict
input = sys.stdin.readline

T = int(input().strip())

n = int(input().strip())
A = [0] + [*map(int, input().strip().split())]
pre_sum_of_A = [0] * (n + 1)

m = int(input().strip())
B = [0] + [*map(int, input().strip().split())]
pre_sum_of_B = [0] * (m + 1)

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

for i in range(1, m + 1):
    pre_sum_of_B[i] = pre_sum_of_B[i - 1] + B[i]

result = 0

# key : A의 부 배열의 합
# value : 해당 합을 갖는 부 배열의 경우의 수
count_sub_sum_of_A = defaultdict(int)

# 위 딕셔너리 채우기
for size in range(1, n + 1):
    for left_idx in range(1, n - size + 2):
        right_idx = left_idx + size - 1
        sub_sum = pre_sum_of_A[right_idx] - pre_sum_of_A[left_idx - 1]

        count_sub_sum_of_A[sub_sum] += 1

# B의 모든 부 배열을 순회하면서, 현재 단계 부 배열의 합이 p라고 할 때,
# A의 부 배열 중 합이 T - p 인 것이 q개이면,
# result에 q를 더해준다.
for size in range(1, m + 1):
    for left_idx in range(1, m - size + 2):
        right_idx = left_idx + size - 1
        sub_sum = pre_sum_of_B[right_idx] - pre_sum_of_B[left_idx - 1]

        result += count_sub_sum_of_A[T - sub_sum]

print(result)

풀이 요약

  1. A의 모든 부 배열을 순회하면서, key: 부 배열의 합, value : 그 합의 경우의 수 딕셔너리를 만들어준다.

  2. B의 모든 부 배열을 순회하면서, 현재 단계의 부 배열의 합이 p라고 할 때, A에서 부 배열의 합이 (T - p) 인 경우의 수가 곧 정답에 더해야할 경우의 수가 된다. 이걸 만들어 둔 딕셔너리를 활용하여 간편하게 result에 더해준다.


어려웠던 점, 배운 점

  1. 내 풀이는 A의 모든 부 배열을 순회하면서, 각 단계에서 부 배열 합이 p라고 할 때, B를 투 포인터로 탐색하여 부 배열 합이 T - p 인 경우를 찾아내는 방식으로 구현했는데, 부 배열 합이 T - p 일 때, 바로 직후에 포인터를 어떻게 움직여야할지에서 막혀서, 결국 다른 사람 풀이를 찾아봤다.

    애초에 내 풀이는 O(n^3) 에 가까운 풀이라 TLE가 예상되었으니 제대로 된 풀이가 아니었을 것 같다. 암튼 다른 사람 풀이의 저 아이디어를 생각해내는게 어려웠다. 어느 한 로직에 매몰되지 않고 폭넓은 시야로 이것저것 아이디어를 생각해볼 수 있는 능력을 길러야겠다.

profile
알고리즘, 풀스택, 앱 개발, 각종 프로젝트 내용 정리 (https://github.com/minsu-cnu)

0개의 댓글

관련 채용 정보