[BOJ] 합이 0인 네 정수 (no.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이다.

🤔 생각

• 오 itertools의 product를 이용하면 뚝딱이겠군. 싶었는...데

• 그걸로 되면 왜 골드 티어인지 생각을 했어야 했다.

• 전체 곱집합을 구하면 반복 돌리는 횟수가 어마어마해진다.

• 그럼 반반씩 나눠서 딕셔너리화 시킨 다음, 각 횟수를 곱하자.

나중에 알았는데 이게 이분 탐색이란다.
그냥 반반 곱집합이지 알고리즘에서 말하는 이분 탐색이랑은 좀 경우가 다르지 않나..?

📌 내 풀이

try-except 문은 선형탐색을 막음으로써 시간복잡도를 단축하기 위해 사용했다.

import sys
input = sys.stdin.readline
print = sys.stdout.write

def get_hash(X, Y):
    hash_table = {}
    for x in X:
        for y in Y:
            part = x + y
            try: hash_table[part] += 1
            except: hash_table[part] = 1

    return hash_table

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_dict = get_hash(A,B)
cd_dict = get_hash(C,D)

answer = 0
for k,v in ab_dict.items():
    try: answer += v * cd_dict[-k]
    except: continue

print("%d" % answer)

생각해보니 굳이 딕셔너리를 두개씩이나 만들 필요 없이 구현할 수 있었다.
2*3 을 할게 아니라, 2+2+2 를 하면 되는거였다.
굳이 횟수를 다 구해둘 필요 없이 그때그때 덧셈해주면 딕셔너리 한개로 가능하다. (공간복잡도를 절반 가까이 단축할 수 있다.)

📌 개선 풀이

import sys
input = sys.stdin.readline
print = sys.stdout.write

n = int(input())
A, B, C, D = [], [], [], []
answer = 0

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

hash_table = {}
for a in A:
    for b in B:
        part = a + b
        try: hash_table[part] += 1
        except: hash_table[part] = 1

for c in C:
    for d in D:
        try: answer += hash_table[-(c+d)]
        except: continue

print("%d" % answer)

✔ 배운 점

• 아직도 의문이다. 이런 것도 이분 탐색이라고 하는건가?
• 같은 이중 for문을 돌릴 때, product를 이용한 쪽이 크게 느리다.
• 쌍을 구할 땐 딕셔너리 곱보다 그때그때 덧셈해서 답을 도출하자.