Coding Test/BaekJoon_Python

백준 2143 <두 배열의 합> Python

JunOnJuly 2024. 9. 8. 04:22
728x90

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


부분합 문제입니다. 단순히 부분합 리스트를 만들고 부분합 / 부분합 쌍을 찾으면 간단히 생각해도 최악의 경우 1000개의 데이터에서 4중 반복문을 돌아야 하기 때문에 다른 방법을 찾아야 합니다.

 

우선 각 부분합 리스트에서 2중 반복문을 통해 모든 부 배열 합을 기록해줍니다. 데이터 기록을 할 때 딕셔너리와 리스트 중 선택할 수 있는데, 딕셔너리는 일반적으로 시간을 위해 공간을 포기한 경우로 해당 문제에서 딕셔너리에 모든 부 배열 합을 기록하면 메모리초과가 발생합니다. 그러므로 리스트에 기록하도록 합니다.

 

기록 중 이분 탐색을 통한 빠른 검색을 위해 정렬을 해주어야 하는데, 이분탐색을 활용한 정렬인 bisect.insort 함수를 사용할 수 도 있지만 이는 데이터를 추가할 때 마다 탐색-정렬을 반복하므로 여러 번 반복되는 작업에서는 모든 데이터를 삽입한 뒤 정렬하는 것이 빠르므로 마지막에 한번에 정렬하도록 합니다.

 

정렬된 A 의 부 배열 합 리스트를 순회하면서 하나의 부 배열 합을 선택합니다. 선택된 부 배열 합은 여러 개 일수 있습니다. 그러므로 반복 작업을 피하기 위해 bisect_right 함수를 사용해 해당 수의 마지막 인덱스를 찾아 해당 합의 개수를 찾습니다. 또 후에 순회하는 인덱스를 합의 개수만큼 뒤로 밀어줍니다.

 

후에 정렬된 B 의 부 배열 합 리스트 중 (T-선택된 A의 부 배열 합) 의 수를 찾아야 합니다. 여기서는 bisect_left, bisect_right 함수들을 사용해 해당 수의 시작과 끝을 골라 해당 합의 개수를 찾습니다.

 

각 합의 수를 구했으므로 이들의 조합(곱)을 통해 개수를 구할 수 있습니다.


import sys
from bisect import bisect_left, bisect_right

input = sys.stdin.readline


def solution(T, n, m, A, B):
    # 결과
    result = 0
    # 부분합 리스트로 치환
    A = [0] + A
    for i in range(1, n+1):
        A[i] += A[i-1]
    B = [0] + B
    for i in range(1, m+1):
        B[i] += B[i-1]
    # 모든 부분합 기록
    # bisect.insort 를 사용하면 쉽게 정렬할 수 있지만
    # 여러번 사용하면 효율이 떨어짐
    # 그냥 다 넣고 마지막에 정렬
    A_subsums = []
    for i in range(0, n):
        for j in range(i+1, n+1):
            A_subsums.append(A[j]-A[i])
    A_subsums.sort()
    B_subsums = []
    for i in range(0, m):
        for j in range(i+1, m+1):
            B_subsums.append(B[j]-B[i])
    B_subsums.sort()
    # A_subsum 을 탐색하며 A 부배열 합 정하기
    i = 0
    while i < len(A_subsums):
        # A 부배열 합
        A_sumsub = A_subsums[i]
        # A 부배열 합과 같은 수 다음 인덱스 검색
        same_A_sumsub = bisect_right(A_subsums, A_sumsub)
        # A 부배열 합과 같은 수 개수
        num_same_A_sumsub = same_A_sumsub - i
        # B 부배열 합
        B_sumsub = T-A_sumsub
        # B 부배열 합과 같은 수 검색
        same_B_sumsub_start = bisect_left(B_subsums, B_sumsub)
        same_B_sumsub_end = bisect_right(B_subsums, B_sumsub)
        num_same_B_sumsub = same_B_sumsub_end - same_B_sumsub_start
        # 결과에 더하기
        result += num_same_A_sumsub * num_same_B_sumsub    
        # i 를 뒤로 이동
        i = same_A_sumsub

    return result


# 입력
T = int(input().strip())
n = int(input().strip())
A = list(map(int, input().strip().split()))
m = int(input().strip())
B = list(map(int, input().strip().split()))

print(solution(T, n, m, A, B))

 

728x90