[BOJ] 7453 - 합이 0인 네 정수 (Python)

 

 

아이디어 #1

 

두 배열씩 묶어 하나의 배열로 총 두 개의 배열을 만든 후, 투 포인터로 합친 배열 값들의 합이 0인 경우를 찾는다.

 

 

풀이 #1

 

입력 값들을 리스트 A, B, C, D에 저장하고, A와 B를 합친 리스트 AB, C와 D를 합친 리스트 CD를 만들어 정렬한다.

투 포인터로 두 리스트 AB와 CD의 원소 합이 0인 경우를 찾고, 중복 값이 있는 경우 그 경우의 수를 구해 결과값에 더한다.

Python 3로는 시간 초과가 나고, PyPy3로 제출해서 통과할 수 있었다.

 

import sys

input = sys.stdin.readline

N = int(input())
A, B, C, D = [], [], [], []
for _ in range(N):
    a, b, c, d = map(int, input().split())
    A.append(a)
    B.append(b)
    C.append(c)
    D.append(d)
A.sort()
B.sort()
C.sort(reverse=True)
D.sort(reverse=True)
AB = [a+b for a in A for b in B]
CD = [c+d for c in C for d in D]
AB.sort()
CD.sort(reverse=True)
M = N*N
i = j = res = 0
while i < M and j < M:
    temp = AB[i]+CD[j]
    if temp == 0:
        cnt_ab = cnt_cd = 1
        while i+1 < M and AB[i] == AB[i+1]:
            i += 1
            cnt_ab += 1
        while j+1 < M and CD[j] == CD[j+1]:
            j += 1
            cnt_cd += 1
        res += cnt_ab*cnt_cd
    elif temp > 0:
        j += 1
        continue
    i += 1
print(res)

 

 

두 리스트를 정렬하지 않은 상태에서 합치고 정렬하는 것보다 정렬한 상태에서 합치고 정렬하는 것이 약 1.8배 더 빨랐다.

# 더 느린 코드
AB = [a+b for a in A for b in B]
CD = [c+d for c in C for d in D]
AB.sort()
CD.sort(reverse=True)

# 더 빠른 코드
A.sort()
B.sort()
C.sort(reverse=True)
D.sort(reverse=True)
AB = [a+b for a in A for b in B]
CD = [c+d for c in C for d in D]
AB.sort()
CD.sort(reverse=True)

 

 

 

아이디어 #2

 

다른 분의 코드를 참고한 결과, 딕셔너리를 이용하면 Python 3로도 간당간당하게 통과할 수 있었다.

 

 

풀이 #2

 

리스트 A와 B의 각 원소들의 합 개수를 딕셔너리 AB에 저장한다.

리스트 C와 D의 각 원소들의 합 중 AB의 원소와의 합이 0이 되는 경우를 결과값에 더해 나간다.

 

같은 코드를 제출하는데도 종종 통과하지 못할 때가 있었고, PyPy3에서는 풀이 #1의 코드보다 약 2.7배 느렸다.

 

import sys

input = sys.stdin.readline

def find():
    AB = dict()
    for a in A:
        for b in B:
            AB[a+b] = AB.get(a+b, 0)+1
    res = 0
    for c in C:
        for d in D:
            res += AB.get(-(c + d), 0)
    return res

N = int(input())
A, B, C, D = [], [], [], []
for _ in range(N):
    a, b, c, d = map(int, input().split())
    A.append(a)
    B.append(b)
    C.append(c)
    D.append(d)
print(find())