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

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

7453번: 합이 0인 네 정수

# 시간제한이 12초인 문제는 처음본다..

 

I. 처음에 A, B배열의 모든 원소의 합을 AB배열에 넣고, C, D배열의 모든 원소의 합을 CD배열에 넣은 뒤

   정렬 한뒤 투 포인터를 사용해 합이 0인 경우를 count해주려고 했지만,

    AB = [0, 0, 0, 0], CD = [0, 0, 0, 0]인 경우 합이 0인 경우 두 포인터 중 하나가 무조건 이동하므로

    또 0이 나오는 경우를 count할 수 없었다.

=====================================================================

 

II. 그래서 CD배열을 정렬한 뒤, bisect모듈의 bisect_right(upper bound)와 bisect_left(lower bound)를

     사용해 AB의 원소에 -를 붙인 값을 찾아서 count를 해주었다.

     이건, AB의 원소의 개수가 N^2개이고 각각의 원소마다 이분탐색을 돌리므로 O(N^2logN^2)이

     나와서 그런지 모르겠지만 시간초과가 발생하였다.

 

import sys, bisect
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)

AB, CD = [], []
for i in range(n):
    for j in range(n):
        AB.append(A[i] + B[j])
        CD.append(C[i] + D[j])
CD.sort()

answer = 0
for i in AB:
    answer += bisect.bisect_right(CD, -i) - bisect.bisect_left(CD, -i)
print(answer)

==========================================================================

 

III. 이분탐색을 사용하지 않고, defaultdict을 이용해 

     A와 B의 모든 합의 경우를 AB라는 딕셔너리에 넣고, C와 D의 모든 합을 CD에 넣고

     AB의 키 값을 반복문으로 돌며 그 키값에 -붙인 값이 CD에 있으면 각각의 개수를 곱해서

     answer에 저장하는 식으로 짰지만, 역시 시간초과가 발생하였다..

      (dict에서 in함수가 평균 O(1)이지만, 최악은 O(N)이어서 시간초과가 나는 것 같다.)

 

AB, CD = defaultdict(int), defaultdict(int)
for i in range(n):
    for j in range(n):
        AB[A[i] + B[j]] += 1
        CD[C[i] + D[j]] += 1

answer = 0
for i in AB:
    if -i in CD:
        answer += AB[i] * CD[-i]
print(answer)

=======================================================================

 

IV. 질문을 해서 어떤분이 답변을 달아주셔서 그 방법으로 했는데 드디어 통과가 되었다!

     AB만 딕셔너리에 넣고, C와 D의 합을 구하고 그 자리에서 바로 AB에 -(C+D)가 있는지

     확인해서 answer에 넣어주는 식으로 짰다. 

 

AB = defaultdict(int)
for i in range(n):
    for j in range(n):
        AB[A[i] + B[j]] += 1

answer = 0
for i in range(n):
    for j in range(n):
        s = C[i] + D[j]
        if -s in AB:
            answer += AB[-s]
print(answer)

(solved.ac 티어: 골드 2)

 

 

# 이렇게 줄이고 pypy3로 제출해도 11.5초로 겨우 시간제한을 통과했다.

# 맞은 사람 3000명 중에 Python3로 통과한 사람이 단 한명도 없다 ㄷㄷ