[BOJ] 7453 합이 0인 네 정수 바로가기
정수로 이루어진 크기가 같은 배열 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이 되는 쌍의 개수를 출력한다.
Two Pointer
) 알고리즘을 적용하여 문제를 해결하였다.Binary Search
) 혹은 투 포인터(Two Pointer
) 알고리즘을 고려해야 한다는 동료의 코멘트가 도움이 되었다.AB, CD = [], []
for i in range(N):
for j in range(N):
AB.append(arr[i][0] + arr[j][1])
CD.append(arr[i][2] + arr[j][3])
AB.sort()
CD.sort(reverse = True)
N * N * N * N
의 모든 경우의 수를 구하기엔 시간이 너무 많이 걸린다.A
, B
원소들로 조합 가능한 모든 조합의 합을 리스트 AB
에 저장한다.C
, D
원소들로 조합 가능한 모든 조합의 합을 리스트 CD
에 저장한다.AB
, CD
리스트를 각각 오름차순, 내림차순으로 정렬한다.while p1 < length and p2 < length:
v1, v2 = AB[p1], CD[p2]
if v1 + v2 == 0:
n1, n2 = 0, 0
while v1 == AB[p1] and p1 < length:
p1, n1 = p1+1, n1+1
if p1 == length: break
while v2 == CD[p2] and p2 < length:
p2, n2 = p2+1, n2+1
if p2 == length: break
answer += n1 * n2
elif v1 + v2 > 0:
p2 += 1
elif v1 + v2 < 0:
p1 += 1
answer
에 누적으로 합한다.from sys import stdin
def solution(N, arr):
answer = 0
# [1] 4개의 리스트 2개로 병합
AB, CD = [], []
for i in range(N):
for j in range(N):
AB.append(arr[i][0] + arr[j][1])
CD.append(arr[i][2] + arr[j][3])
AB.sort()
CD.sort(reverse = True)
p1, p2 = 0, 0
length = N*N
# [2] 투포인터 체크
while p1 < length and p2 < length:
v1, v2 = AB[p1], CD[p2]
if v1 + v2 == 0:
n1, n2 = 0, 0
while v1 == AB[p1] and p1 < length:
p1, n1 = p1+1, n1+1
if p1 == length: break
while v2 == CD[p2] and p2 < length:
p2, n2 = p2+1, n2+1
if p2 == length: break
answer += n1 * n2
elif v1 + v2 > 0:
p2 += 1
elif v1 + v2 < 0:
p1 += 1
return answer
# input
N = int(stdin.readline())
arr = list(list(map(int,stdin.readline().split())) for _ in range(N))
# result
result = solution(N, arr)
print(result)