Baekjoon 2143. 두 배열의 합 / Python

2143. 두 배열의 합

난이도 : 골드 3
소요 시간 : 25분
날짜 : 2024.12.23
언어 : 파이썬
알고리즘 유형 : 이분탐색, 누적

설명 보기전에 문제 풀어보러 가기

1. 문제 설명

1. 정수 t와 배열 a, 배열 b 가 주어진다.
2. 배열 a와 배열 b에서 각각 모든 원소가 연속된 배열을 선택한다.
3. 선택된 배열들의 숫자의 합이 t 가 되는 경우의 수를 구한다.

2. 해결 방식

1. 배열a 와 배열 b 에서 나올 수 있는 모든 합을 구해서 딕셔너리형태로 저장했다.
2. 배열a 를 순회하여 그 값(key)과 개수(value)를 구하고 배열 b에서 t - key의 value(개수)를 찾아서 곱한 후, 결과에 누적한다.

3. 정답 코드

import sys
sys.stdin = open("C:/Users/ghtjd/Desktop/tmp/python/input.txt", "r")

# input = sys.stdin.readline

from collections import defaultdict

t = int(input())

n = int(input())
a = list(map(int, input().split()))

m = int(input())
b = list(map(int, input().split()))

# 부분집합의 합 구해놓기, 단 인덱스가 연속이어야 함
def subset_sum(N, arr):
    dict_arr = defaultdict(int)
    # 크기가 x 인 부분집합
    for x in range(1, N + 1):
        # 누적합 알고리즘
        s = sum(arr[:x])
        dict_arr[s] += 1
        for i in range(1, N - x + 1):
            s += arr[i + x - 1] - arr[i - 1]
            dict_arr[s] += 1 

    return dict_arr

dict_A = subset_sum(n, a)
dict_B = subset_sum(m, b)

res = 0

for k_A, v_A in dict_A.items():
    res += v_A * dict_B[t - k_A]

print(res)