[백준] 7453번: 합이 0인 네 정수 (Python3)

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

 

문제

정수로 이루어진 크기가 같은 배열 A, B, C, D가 있다.

A[a], B[b], C[c], D[d]의 합이 0인 (a, b, c, d) 쌍의 개수를 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 배열의 크기 n (1 ≤ n ≤ 4000)이 주어진다. 다음 n개 줄에는 A, B, C, D에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 228이다.

 

출력

합이 0이 되는 쌍의 개수를 출력한다.

 

예제 입력 1

6
-45 22 42 -16
-41 -27 56 30
-36 53 -37 77
-36 30 -75 -46
26 -38 -10 62
-32 -54 -6 45

 

예제 출력 1

5

---

문제 설명

n 이 입력되고, n번 동안 각 List A, B, C, D의 element가 하나씩 입력됩니다.

그럼 이 list 에서 각 List A, B, C, D에서 하나의 element씩을 뽑아서

그 합이 0으로 만들 수 있는 순열은 몇 개가 있는지 결정하면 됩니다.

 

전략

우선 모든 경우의 수를 탐색하면 len(A) * len(B) * len(C) * len(D) 의 iteration이 발생하므로,

시간이 길게 주어졌지만 O(n^4)으로는 문제를 시간 내에 해결할 가능성이 없을 것 입니다.

 

핵심 전략은 A와 B 경우의 수를 계산, C와 D의 경우의 수를 계산하여

이 둘을 비교하여 n^2 iteration 두 번으로 해결하는 것입니다.

A와 B의 조합으로 가능한 모든 경우의 수를 Dictionary 형태로

key = 가능한 수, value = 경우의 수 로 만드는 것입니다.

예를 들어, A, B의 element의 조합에 의해 43이 3번 만들어진다면

x[43] = 3 이런식으로 말입니다.

 

그리고 C와 D의 조합으로 가능한 모든 경우의 수에 대해

그 합이 0이 되는 수가 dictionary에 존재하는지를 판단합니다.

예를 들어, C와 D의 조합으로 -43이 만들어졌다면,

위에서 만들어둔 Dictionary x에서 x[43] 을 찾고, 

그 경우의 수만큼을 결과에 더해주는 것입니다.

 

코드

#7453
import sys
input = sys.stdin.readline

n = int(input())

A = []
B = []
C = []
D = []

# A, B, C, D list 생성
for _ in range(n):
    a, b, c, d = map(int, input().split())
    A.append(a)
    B.append(b) 
    C.append(c)
    D.append(d)   

x = dict() # dictionary 생성

# A, B 조합 경우의 수 고려, dict 작성
for a in A:
    for b in B:
        if a+b in x:
            x[a+b] += 1
        else:
            x[a+b] = 1

result = 0 # 합 0이 되는 경우의 수

# C, D 조합 경우의 수 고려, dict에서 합이 0이 되는 수를 결과에 합산
for c in C:
    for d in D:
        if -(c+d) in x:
            result += x[-(c+d)]

print(result)