[백준 7456] 합이 0인 네 정수 (Python)

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

7453번: 합이 0인 네 정수

 

나의 풀이(시간 초과)

 

import sys
input = sys.stdin.readline
n = int(input())
A, B, C, D = [],[],[],[]
for _ in range(n):
    q,w,e,r = map(int,input().split())
    A.append(q)
    B.append(w)
    C.append(e)
    D.append(r)

A = sorted(A)
B = sorted(B)
C = sorted(C)
D = sorted(D)
cnt = 0
for i in range(n):
    a = A[i]
    for j in range(n):
        b = B[j]
        for k in range(n):
            c = C[k]
            for l in range(n):
                d = D[l]
                if sum([a,b,c,d]) == 0:
                    temp.append([a, b, c, d])
                    cnt += 1
print(cnt)

 

당연히 시간초과가 날것이라고 생각했고 도저히 아이디어가 떠오르지 않아 풀이를 찾게 되었다.

 

재도전(실패)

import sys
input = sys.stdin.readline
n = int(input())
A, B, C, D = [],[],[],[]
arr = [list(map(int,input().split())) for _ in range(n)]
ab = []
cd = []
for i in range(n):
    for j in range(n):
        ab.append(arr[i][0] + arr[j][1])
        cd.append(arr[i][2] + arr[j][3])
ab = list(set(ab))
cd = list(set(cd))
ab.sort()
cd.sort(reverse = True)
ab_start, cd_start = 0, 0
cnt = 0
while ab_start < len(ab) and cd_start < len(cd):
    if ab_start >= len(ab) or (cd_start < len(cd) and ab[ab_start] + cd[cd_start] < 0):
        ab_start += 1
    else:
        if ab[ab_start] + cd[cd_start] == 0:
            cnt += 1
        cd_start += 1
print(cnt)

시간초과가 나버렸다

 

찾아본 정답 코드

 

import sys
input = sys.stdin.readline

if __name__ == "__main__":
    n = int(input())
    arr = [list(map(int, input().split())) for _ in range(n)]
    ab, cd = [], []
    for i in range(n):
        for j in range(n):
            ab.append(arr[i][0] + arr[j][1])
            cd.append(arr[i][2] + arr[j][3])
    
    ab.sort()
    cd.sort()
    # print(ab, cd)
    i, j = 0, len(cd) - 1
    result = 0
    while i < len(ab) and j >= 0:
        if ab[i] + cd[j] == 0: 
            nexti, nextj = i + 1, j - 1

            while nexti < len(ab) and ab[i] == ab[nexti]: 
                nexti += 1

            while nextj >= 0 and cd[j] == cd[nextj]: 
                nextj -= 1
            
            result += (nexti - i) * (j - nextj)
            i, j = nexti, nextj

        elif ab[i] + cd[j] > 0:
            j -= 1
        else:
            i += 1

    print(result)

https://velog.io/@ckstn0778/%EB%B0%B1%EC%A4%80-7453%EB%B2%88-%ED%95%A9%EC%9D%B4-0%EC%9D%B8-%EB%84%A4-%EC%A0%95%EC%88%98-X-1

백준 7453번 - 합이 0인 네 정수(★★★ / X▲ / 2) : Python