7453. 합이 0인 네 정수 (Python)

7453번: 합이 0인 네 정수

 

풀이

 두 가지 방식으로 풀었는데 기본적으로 적용되는 것은 A, B, C, D 이렇게 4개의 배열을 두고 각각 더하는 것이 아니라, A 배열과 B 배열의 원소를 합한 결과(이하 AB라고 표기) 그리고 C 배열과 D 배열의 원소를 합한 결과(이하 CD라고 표기) 이렇게 두 개로 나누어 구한다. 이때 AB의 원소와 CD의 원소의 합이 0이 되려면 두 값의 부호가 반대여야 한다(ex. AB: 29, CD: -29).

 

(1) Dict를 사용한 풀이

 먼저 A와 B를 합한 값을 키로 하여 딕셔너리에 저장한다. 그 후 C와 D를 합하면서 해당 값의 부호를 바꾼 값이 AB 딕셔너리에 존재한다면 해당 키의 값을 결과 변수에 더해주면 된다.

 처음에는 AB 딕셔너리를 만들고 CD 딕셔너리를 만든 후 for문을 통해 AB의 키를 탐색하면서 해당 키에 -1을 곱한 값을 CD가 키로 갖고 있는지를 체크하고 이를 결과 변수에 더해주는 방식으로 구현했으나 시간 초과가 발생했다. 시간초과를 피하기 위해서는 CD 딕셔너리를 따로 만들지 말고 for문을 통해 C와 D를 더하는 즉시 부호를 바꾼 값이 AB 딕셔너리에 존재하는지를 체크 해야 한다.

 

import sys
input = sys.stdin.readline

if __name__ == '__main__':
    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)

    abDict = {} # A배열과 B배열을 더한 값을 키로 갖는 딕셔너리 
    for a in A:
        for b in B:
            abDict.setdefault(a+b, 0) 
            abDict[a+b] += 1

    # C와 D를 더하면서 abDict에서 -(c+d)를 키로 갖고 있다면 값을 가져와 res에 더해준다.
    res = 0
    for c in C:
        for d in D:
            res += abDict.get(-(c+d), 0)

    print(res)

 

(2) 투 포인터와 이분 탐색을 사용한 풀이

 다음은 투 포인터와 이분 탐색을 사용한 풀이다. 이번엔 AB와 CD를 리스트로 둔다. 그 후 AB와 CD를 모두 오름차순으로 정렬한다. 이제 포인터를 두는데 left 포인터는 AB의 맨 앞에 두고 right 포인터는 CD의 맨 뒤에 둔다. 합해서 0이 되려면 AB와 CD의 값의 부호가 달라야 하기 때문이다. 

 투 포인터를 진행하는 과정은 다음과 같다.

1. AB[left] + CD[right]가 0보다 작은 경우
   -> 0이 되기 위해서는 두 수의 합이 커져야 한다. 합을 크게 하기 위해서 left를 오른쪽으로 1 이동시킨다. 
2. AB[left] + CD[right]가 0보다 큰 경우
   -> 0이 되기 위해서는 두 수의 합이 작아져야 한다. 따라서 right에 -1을 해주어 큰 쪽의 수를 줄여야 한다.
   
   
3. AB[left] + CD[right]가 0일 경우
   -> 이 때는 단순히 카운트를 1 증가시키면 안된다. 이유는 AB나 CD에 같은 값이 여러 개 있을 수 있기 때문이다. 그렇기 때문에 AB[left]의 값이 AB에 몇 개가 있고 CD[right]의 값이 CD에 몇 개가 있는지를 체크해야 한다. 이를 위해 이분 탐색을 사용하여 upper bound와 lower bound를 구해야 한다. upper bound - lower bound를 하면 같은 값이 해당 리스트에 몇 개가 있는지를 구할 수 있다.
    파이썬에서는 bisect_left와 bisect_right를 사용한다. 이를 사용하여 카운트 한 후 left는 upper bound를 할당하고 right는 lower bound-1을 할당해주면 새로운 숫자에 대하여 체크할 수 있다.
   
   

import bisect
import sys
input = sys.stdin.readline

if __name__ == '__main__':
    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])
    
    # 두 배열 모두 오름차순으로 정렬 
    ab.sort()
    cd.sort()

    # left는 ab 배열의 시작 부분, right는 cd 배열의 마지막 부분
    left, right = 0, len(cd)-1 
    res = 0

    while left < len(ab) and right >= 0:
        tmp = ab[left] + cd[right]

        if tmp > 0: # 합이 0보다 클 경우 right에 -1을 해주어 합을 줄여준다.
            right -= 1
        elif tmp < 0: # 합이 0보다 작을 경우 left에 +1을 해주어 합을 높여준다.
            left += 1
        else: # 0일 경우에는 각각의 값이 배열에 몇 개 존재하는지를 체크해야 한다.
            # upperBound와 lowerBound를 사용하여 배열에 같은 값이 몇 개 있는지를 구할 수 있다.
            abLeft, abRight = bisect.bisect_left(ab, ab[left]), bisect.bisect_right(ab, ab[left])
            cdLeft, cdRight = bisect.bisect_left(cd, cd[right]), bisect.bisect_right(cd, cd[right])

            # upperBound - lowerBound를 통해서 같은 값이 몇 개 있는지를 구하고
            # 두 값을 곱하여 경우의 수를 카운트한다.
            # 그 후, left는 upperBound로 right는 lowerBound-1로 설정하면
            # 다시 새로운 값부터 합이 0이 되는 경우를 체크할 수 있다.
            res += (abRight - abLeft) * (cdRight - cdLeft)
            left = abRight
            right = cdLeft - 1
    
    print(res)