[BOJ] Q7453 합이 0인 네 정수

[BOJ] Q7453 합이 0인 네 정수

Question

Language: Python

Difficulty: Gold 2

해당 문제는 주어진 배열들에 대해 네 정수의 합이 0이 되는 순서쌍의 갯수를 구하는 문제이다.

우선, 처음 고려해볼 수 있는 것이 모든 순서쌍의 갯수를 비교하는 Bruteforce 방식을 확인해본다. 위 문제의 경우 4개의 배열이 있으므로, 총 O(n⁴)의 시간복잡도가 소요되고, n이 최대 4000개이므로 시간 내에 풀이하는 것이 불가능하다.

따라서, 차원의 갯수를 조금 줄이는 것이 필요하다 즉, 4개의 순서쌍을 2개의 순서쌍을 낮춰서 고려한다. 아래와 같이 2개의 정수합 형태로 묶어서 서로 비교한다.

for i in range(n):
    for j in range(n):
        ab_sum.append(A[i]+B[j])
        cd_sum.append(C[i]+D[j])
    

위 2개의 배열을 활용해서 Two-pointer 방식을 통해 합을 비교하면서 합이 0이 되는 순서쌍의 갯수를 구한다.

Solution 1

from bisect import bisect_left,bisect_right
def solution():
    count=0
    ab_sum=[]
    cd_sum=[]

    for i in range(n):
        for j in range(n):
            ab_sum.append(A[i]+B[j])
            cd_sum.append(C[i]+D[j])
    
    ab_sum.sort()
    cd_sum.sort()
    i,j=0,n**2-1
    while i < n**2 and j >=0:

        if ab_sum[i]+cd_sum[j]==0:
            i+=1
            j-=1
            
            left_count,right_count=1,1
            #왼쪽 배열에서 같은 수의 갯수
            while i<n**2 and ab_sum[i] == ab_sum[i-1]:
                i+=1
                left_count+=1
            #오른쪽 배열에서 같은 수의 갯수
            while j>=0 and cd_sum[j] == cd_sum[j+1]:
                j-=1
                right_count+=1
            #순서쌍의 갯수는 왼쪽 갯수 * 오른쪽 갯수
            count+=(left_count*right_count)
        #만약 합이 0보다 작다는 의미는 음수의 절대값이 큼을 의미하기 때문에 음수 측 절대값을 낮춰야한다.
        elif ab_sum[i] + cd_sum[j] <0:
            i+=1
        else:
            j-=1

    print(count)
            

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)

    solution()

Solution 2

defaultdict을 활용하여 A[a]+B[b]의 값을 저장하고, C[c]+D[d]의 값을 defaultdict에 찾으므로써 순서쌍의 갯수를 구할 수도 있다.

from collections import defaultdict
def solution():
    count=0
    ab_sum=defaultdict(int)

    #
    for i in range(n):
        for j in range(n):
            ab_sum[-(A[i]+B[j])]+=1

    for i in range(n):
        for j in range(n):
            if (C[i]+D[j]) in ab_sum.keys():
                count+=ab_sum[C[i]+D[j]]
    print(count)
            

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)

    solution()