[백준] 7453. 합이 0인 네 정수

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에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 2^28이다.

출력

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

풀이

단계별로 알고리즘을 적용해보자.

완전 탐색

위 문제를 보자마자 생각나는 풀이는 완전 탐색이다.
아래가 완전 탐색 코드다. 그리고 이렇게 하면 백퍼 시간초과에 걸린다.

import sys
input = sys.stdin.readline
n = int(input())
a, b, c, d = list(), list(), list(), list()

for i in range(n):
    t, e, m, p = map(int, input().split())
    a.append(t)
    b.append(e)
    c.append(m)
    d.append(p)

answer = list()

for i in range(n):
    for j in range(n):
        for k in range(n):
            for m in range(n):
                if a[i] + b[j] + c[k] + d[m] == 0:
                    answer.append([i, j, k, m])

print(len(answer))

for loop 4개를 냅다 해버리면 시간복잡도가 O(n^4)가 되고 1~2%에서 시간 초과가 날 것이다.
완전 탐색으로 풀면 안되는 문제이다. (애초에 카테고리에 완전 탐색이 없다)

for loop 줄이기

https://jaeyoon-95.tistory.com/168
참고해서 풀었다.
대부분 dictionary로 시간복잡도 O(n^2)로 푼 것 같다.
배열 a와 b를 먼저 더하고 배열 c와 d를 더해서 문제를 풀었다.

파이썬 dictionary에 get 함수를 잘 사용하면 되었고 이번 기회에 알게 되었다.

get은 keyError가 뜨지 않고 default 값을 지정해주지 않으면 None이 뜨고, default 값을 지정해주면 해당 값이 뜬다고 한다.

위의 함수를 모르고 코드를 짰을 땐 다음과 같이 적었다.

a = dict()
count = 0

if not 'key' in a.keys():
  a[key] = 0
count += a.get(key)

근데 이젠 그냥 다음과 같이 쓰면 위와 같다는 것을 알게 되었당..!

a = dict()
count = 0

count += a.get(key, 0)

풀이는 다음과 같다.

import sys
input = sys.stdin.readline
n = int(input())
a, b, c, d = list(), list(), list(), list()

for i in range(n):
    t, e, m, p = map(int, input().split())
    a.append(t)
    b.append(e)
    c.append(m)
    d.append(p)

answer = dict()
for _a in a:
    for _b in b:
        answer[_a + _b] = answer.get(_a+_b, 0) + 1
count = 0
for _c in c:
    for _d in d:
        count += answer.get(-(_c+_d), 0)
print(count)