백준 2143 - 두 배열의 합 (Python)

https://acmicpc.net/problem/2143 

2143번: 두 배열의 합

원래는 이분 탐색으로 푸는 문제인데, 그냥 dictionary와 누적합으로 풀었다.

이게 왜 됨? 진짜 파이썬은 알다가도 모르겠다.

처음에 생각한 방법으로는 용량이 부족할 것 같아서 안 했는데, 결국에는 처음 생각한 방법으로도 풀리더라.

 

1. DP를 이용해 시도한 코드 - TLE

import sys
input=sys.stdin.readline
t=int(input())
a_len=int(input())
a=list(map(int,input().split()))
dp_a=[[0 for _ in range(a_len)] for _ in range(a_len)]
for i in range(a_len):
    dp_a[i][i]=a[i]
    for j in range(i+1,a_len):
        dp_a[i][j]=dp_a[i][j-1]+a[j]

b_len=int(input())
b=list(map(int,input().split()))
dp_b=[[0 for _ in range(b_len)] for _ in range(b_len)]
for i in range(b_len):
    dp_b[i][i] = b[i]
    for j in range(i+1,b_len):
        dp_b[i][j]=dp_b[i][j-1]+b[j]

answer=0
for i in range(a_len):
    for j in range(i,a_len):
        temp_a=dp_a[i][j]
        for k in range(b_len):
            for l in range(k,b_len):
                temp_b=dp_b[k][l]
                if temp_a+temp_b==t:
                    answer+=1

print(answer)

당연히 TLE. 생각나는대로 해 봤는데 되지 않았다.

 

2. 누적 합과 dictionary를 이용한 풀이

# 두 배열의 합
# 더 긴 배열의 누적합 값을 다 hash해둔다.
# 더 작은 배열의 누적합을 t에서 뺀 값이 dict에 있으면 그 개수만큼 정답에 더하기. 
import sys
input=sys.stdin.readline
t=int(input())
numbers=dict() # 더 긴 배열의 누적합 값들을 저장할 dictionary이다.
answer=0

a_len=int(input())
a=list(map(int,input().split()))
for i in range(1,a_len):
    a[i]=a[i-1]+a[i]
a=[0]+a # 편의상 [0]을 앞에 더한다.

b_len=int(input())
b=list(map(int,input().split()))
for i in range(1,b_len):
    b[i]=b[i-1]+b[i]
b=[0]+b

if a_len<b_len: # 더 긴 것으로 hash 하는 게 이득이다.
    b_len,a_len=a_len,b_len
    b,a=a,b

for end in range(1,a_len+1): # 누적합 따져보기 시작.
    for start in range(end):
        temp=a[end]-a[start]
        if temp not in numbers: # 값이 dict에 없다면 만들고, 있다면 +=1
            numbers[temp]=1
        else:
            numbers[temp]+=1

for end in range(1,b_len+1): 
    for start in range(end):
        temp=b[end]-b[start]
        if t-temp in numbers: # 지금 보는 누적합을 t에서 뺀 것이 dict에 있다면 answer+=dict[t-누적합]
            answer+=numbers[t-temp]

print(answer)

메모리 초과가 날 줄 알았는데 아니었다.

 

3. 이분 탐색을 이용한 풀이

써 보지는 않았는데, 방법은 찾아서 이해했다. 요는 A의 누적합을 하나의 리스트에 싹다 넣고 sort 한다.

그 후에 B의 누적합을 하나 하나 보면서 이분 탐색을 하며 left_bound와 right_bound를 찾아 두 개를 뺀다.

그러면 개수가 나온다. 다만 그 방법보다는 내가 한 풀이가 더 깔끔하지 않은가 싶었다.