백준 2143 [두 배열의 합] 파이썬

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

 

# 백준 2143 두 배열의 합

import sys

# 목표 합 T 입력
T = int(sys.stdin.readline())

# 첫 번째 배열 A의 크기와 원소 입력
N = int(sys.stdin.readline())
A = list(map(int, sys.stdin.readline().split()))

# 두 번째 배열 B의 크기와 원소 입력
M = int(sys.stdin.readline())
B = list(map(int, sys.stdin.readline().split()))

ans = 0  # 정답 변수 (쌍의 개수)
a_dic = dict()  # A의 부분합 빈도 저장용 딕셔너리
b_dic = dict()  # B의 부분합 빈도 저장용 딕셔너리

# A 배열의 모든 부분합을 계산하여 빈도 저장
for i in range(N):
    for j in range(i, N):
        s = sum(A[i:j+1])  # A[i]부터 A[j]까지 부분합
        if s not in a_dic:
            a_dic[s] = 1
        else:
            a_dic[s] += 1

# B 배열의 모든 부분합을 계산하여 빈도 저장
for i in range(M):
    for j in range(i, M):
        s = sum(B[i:j+1])  # B[i]부터 B[j]까지 부분합
        if s not in b_dic:
            b_dic[s] = 1
        else:
            b_dic[s] += 1

# A의 부분합과 B의 부분합을 합쳐서 T가 되는 경우의 수 계산
for val in a_dic.keys():
    # T - val 값이 B의 부분합에 존재하면, 가능한 조합의 수를 곱해서 누적
    if (T - val) in b_dic:
        ans += (a_dic[val] * b_dic[T - val])

# 결과 출력
print(ans)

합을 딕셔너리에 저장하여 해결하였다. (딕셔너리 요소 탐색 시 시간복잡도 O(1))

 

import sys

# 입력 받기
T = int(sys.stdin.readline())
N = int(sys.stdin.readline())
A = list(map(int, sys.stdin.readline().split()))
M = int(sys.stdin.readline())
B = list(map(int, sys.stdin.readline().split()))

# A의 모든 부분합 구하기
subA = []
for i in range(N):
    s = 0
    for j in range(i, N):
        s += A[j]
        subA.append(s)

# B의 모든 부분합 구하기
subB = []
for i in range(M):
    s = 0
    for j in range(i, M):
        s += B[j]
        subB.append(s)

# 부분합 리스트 정렬 (투포인터를 위해)
subA.sort()
subB.sort(reverse=True)  # 하나는 오름차순, 하나는 내림차순

# 투 포인터 초기화
p1 = 0
p2 = 0
ans = 0

# subA는 오름차순, subB는 내림차순
while p1 < len(subA) and p2 < len(subB):
    sum_val = subA[p1] + subB[p2]
    
    if sum_val == T:
        # 같은 값이 몇 개씩 있는지 세기
        a_val = subA[p1]
        b_val = subB[p2]
        cnt_a = 0
        cnt_b = 0
        
        while p1 < len(subA) and subA[p1] == a_val:
            cnt_a += 1
            p1 += 1
        while p2 < len(subB) and subB[p2] == b_val:
            cnt_b += 1
            p2 += 1

        ans += cnt_a * cnt_b  # 가능한 쌍의 개수 누적

    elif sum_val < T:
        p1 += 1
    else:
        p2 += 1

# 결과 출력
print(ans)

투 포인터로 해결한 코드도 첨부한다.