💡 문제
정수로 이루어진 크기가 같은 배열 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이 되는 쌍의 개수를 출력한다.
💭 접근
이 문제를 해결하기 위해 각각의 배열에서 4중 for문으로 원소를 하나씩 바꿔가며 합을 구하는 단순한 방식으로 구현할 시 입력의 최대가 4000이므로 최악의 경우 4000 만큼의 경우의 수를 보아야 한다.
이렇게 할 경우 당연히 시간초과가 발생하게 되고, 이를 방지하기 위해 1208번 부분수열의 합 2 문제와 유사하게 해결할 수 있다.
우선, 배열 A, B, C, D를 한번에 고려하는 것이 아닌 A와 B, C와 D로 나누어 고려해주면 된다.
이를 위해 먼저 배열 A, B에서 2중 for문을 사용하여 원소를 뽑아 이들의 합을 key로 가지는 딕셔너리에 경우의 수를 저장해준다.
cnt = dict()
for i in range(n):
for j in range(n):
if a_arr[i] + b_arr[j] not in cnt.keys():
cnt[a_arr[i] + b_arr[j]] = 1
else:
cnt[a_arr[i] + b_arr[j]] += 1예를 들어 배열 A = [1, 2], B = [1, 2]가 들어있다고 가정했을때 나올 수 있는 수의 합은 [2, 3, 3, 4]가 되고 이를 딕셔너리에 저장하면 다음과 같다.
cnt = {2: 1,
3: 2,
4: 1}이후, 배열 C와 D를 고려해줄 차례인데 네 배열의 각각의 원소의 합이 0이되는 경우의 수를 구해야 하므로 위 경우에 비추어보면, C와 D의 합은 [-2, -3, -3, -4]가 되어야 한다.
ans = 0
for i in range(n):
for j in range(n):
if -(c_arr[i] + d_arr[j]) not in cnt.keys():
pass
else:
ans += cnt[-(c_arr[i] + d_arr[j])]이처럼 C와 D의 합에 -1을 곱한 값을 key로 가지는 경우의 수를 정답에 더해주면 해결할 수 있다.
📒 코드
n = int(input())
a_arr = []
b_arr = []
c_arr = []
d_arr = []
for _ in range(n):
a, b, c, d = map(int, input().split())
a_arr.append(a)
b_arr.append(b)
c_arr.append(c)
d_arr.append(d)
cnt = dict()
for i in range(n):
for j in range(n):
if a_arr[i] + b_arr[j] not in cnt.keys():
cnt[a_arr[i] + b_arr[j]] = 1
else:
cnt[a_arr[i] + b_arr[j]] += 1
ans = 0
for i in range(n):
for j in range(n):
if -(c_arr[i] + d_arr[j]) not in cnt.keys():
pass
else:
ans += cnt[-(c_arr[i] + d_arr[j])]
print(ans)💭 후기
제한시간이 12초라서 당황했지만 생각보다 쉽게 해결한 문제였다.
🔗 문제 출처
https://www.acmicpc.net/problem/7453
