[백준/7453] 합이 0인 네 정수 (Python)

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

 

Difficulty : Gold 2

 

Status : Solved(pypy3)

 

Time : 01:18:03

 

---

 

문제 설명

 

 

정수로 이루어진 크기가 같은 배열 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이 되는 쌍의 개수를 출력한다.

 

입력 예시

 

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

 

출력 예시

 

5

 

 

---

 

풀이

 

이분 탐색으로 개고생. Defaultdict은 자체 처리 시간이 길다는 교훈 하나를 얻고 투 포인터로 시도해보았다.

 

A + B + C + D를 (A+B) + (C+D)로 묶어 줄 수 있으므로, 가능한 경우의 수를 따져 두 개의 리스트로 바꿔주자(O(N^2)), 이를 정렬하고, 투 포인터로 합이 되는 경우의 수를 찾아주면 된다. (O(N^2))

 

 

 

풀이 코드

import sys

N = int(input())
num_list = [list() for _ in range(4)]

for _ in range(N) :
  for i, num in enumerate(map(int, input().split())) :
    num_list[i].append(num)

new_num_list = [list() for _ in range(2)]

for i in range(N) :
  for j in range(N) :
    for k in range(2) :
      new_num_list[k].append(num_list[2*k][i] + num_list[2*k+1][j])

for i in range(2) :
  new_num_list[i].sort()

answer, i, j = 0, 0, N**2-1

while i < N**2 and j > -1 :
  if new_num_list[0][i] + new_num_list[1][j] == 0 :
    ti, tj = i, j
    while i < N**2 and new_num_list[0][ti] == new_num_list[0][i] :
      i += 1
    while j > -1 and new_num_list[1][tj] == new_num_list[1][j] :
      j -= 1
    answer += (i - ti) * (tj - j)
  elif new_num_list[0][i] + new_num_list[1][j] > 0 :
    j -= 1
  else :
    i += 1
    
print(answer)

풀이 완료!