7453번 - 합이 0인 네 정수(python3)

문제 링크: www.acmicpc.net/problem/7453

7453번: 합이 0인 네 정수

A, B, C, D 4개의 list에 n개의 숫자가 들어있습니다. 각 list에서 숫자 한 개씩 뽑아 합이 0이 되도록 하고 싶습니다. 가능한 총 가짓수를 출력하면 됩니다. 우선, python3으로 제출하면 무조건 시간 초과가 뜨는 것 같으니 pypy3으로 꼭 제출하시기 바랍니다.

 

이전에 본 세 용액 문제처럼 브루트 포스로 접근하면 시간 복잡도가 $O(n^{4})$가 되어 시간 초과가 뜹니다. 따라서, 세 용액 문제에서 접근했던 투 포인터 방식으로 문제를 풀어봅시다. 먼저 전체 코드입니다.

 

세 용액과는 다르게 이번에는 한 번에 더해지는 숫자가 4개입니다. 포인터 2개를 움직이기 위해서는 2개의 숫자를 고정시켜야 하는데, 어떻게 풀 수 있을까요? 먼저 A,B를 묶어 한 개의 리스트 AB_sum으로, C,D를 묶어 한 개의 리스트 CD_sum으로 만들고 정렬시킵니다. 정렬된 두 리스트 AB_sum, CD_sum에 각각 포인터 하나를 배정하고, 움직이면서 0이 되는 모든 가짓수를 찾으면 됩니다. AB_sum에는 A와 B의 원소들의 합으로 만들 수 있는 총 $n^{2}$개의 결과가 저장됩니다. CD_sum도 마찬가지구요. itertools.product를 이용해 빠르게 가능한 모든 [a,b] pair로 구성된 리스트를 만들고 각 pair의 sum을 저장해도 되는데, 어떤 방법이 더 빠른지는 모르겠습니다.

AB_sum = sorted([a+b for a in A for b in B])
CD_sum = sorted([c+d for c in C for d in D])

예시, left pointer와 right pointer의 초기 상태

정렬된 AB_sum, CD_sum을 얻었다면, AB_sum의 첫 번째 원소(AB_sum의 최솟값)를 left pointer(l_p)가 가리키도록, CD_sum의 마지막 원소(CD_sum의 최댓값)를 r_p가 가리키도록 합니다. 두 포인터가 가리키는 값들의 합이 0보다 크면 작아져야 하므로 r_p를 왼쪽으로 1칸, 0보다 작으면 커져야 하므로 l_p를 오른쪽으로 1칸씩 옮깁니다.

 

만약 두 포인터가 가리키는 값들의 합이 0이 된다면? 찾았으므로 가짓수를 올려줘야 합니다. 하지만 바로 sol += 1을 해버리면 문제가 발생합니다. 위 그림과 같은 경우 AB_sum[1] ~ AB_sum[3], CD_sum[3]~CD_sum[4]에서 합이 0이 되는 총 6가지 경우가 발생하는데, 무턱대고 올려버리면 이를 반영하지 못합니다. 따라서 합이 0인 지점을 발견했을 때, 연달아서 같은 수가 몇 번 나오는지를 체크해줘야 합니다.

합이 0이 되는 상태에서 두 포인터의 이동

먼저, l_p가 리스트 인덱스 내에 있다는 전제 하에, 현재 l_p가 가리키는 값과 다음 값이 같은지 확인합니다. AB_sum[1] + CD_sum[4] = 0인데, AB_sum[2] + CD_sum[4] 또한 마찬가지로 0이니까요. l_p가 AB_sum[3]을 가리키는 상태에서는 AB_sum[3] != AB_sum[4]니까 현 상태에서 이동한 거리인 3을 l_move에 저장하고 l_p의 이동을 마칩니다.

 

r_p도 똑같이 움직입니다. r_p의 경우 CD_sum[3]까지 움직이고, r_move에 2를 저장한 상태로 이동을 멈추게 되겠죠. l_p와 r_p 모두 이동을 멈추면, l_move*r_move값인 6을 sol에 더해주고, 마지막으로 l_p와 r_p를 한 칸씩 옮겨주고 마칩니다. 다음부터는 다시 두 포인터가 가리키는 값들의 합이 0보다 큰지, 작은지, 같은지에 따라 이동하는 단계로 돌아갑니다.

  else: #0이면?
      l_move = 1 #연속해서 같은 값이 나올 수 있으므로 작업이 조금 필요함
      r_move = 1
      while True:
          if l_p < len(AB_sum)-1 and AB_sum[l_p] == AB_sum[l_p+1]:
              l_move += 1 #인덱스 범위 안이면서 같은 값이 연달아 나올 경우
              l_p += 1
          elif r_p > 0 and CD_sum[r_p] == CD_sum[r_p-1]:
              r_move += 1 #인덱스 범위 안이면서 같은 값이 연달아 나올 경우
              r_p -= 1
          else: #다음으로 두 포인터가 가리키는 숫자들의 합이 0이 아닌 경우 
              sol += l_move*r_move #결과값 저장
              l_p += 1
              r_p -= 1
              break

처음에 AB_sum, CD_sum을 만드는 데 $O(n^{2})$, 정렬하는데 $O(n^{2}logn^{2})$, 포인터를 돌리는데 $O(n^{2})$이 걸리므로 총 시간복잡도는 $O(n^{2}logn)$인 것 같습니다. 아무튼 브루트 포스보다는 확실히 줄어드네요.

 

*다음은 시간 초과가 떠서 실패한 코드입니다.

 

itertools를 사용했고, dictionary를 사용해 두 포인터가 가리키는 숫자의 합이 0일 때 조건을 확인하면서 포인터를 옮기는 것이 아니라, value값으로 해당 key가 몇 번 등장했는지를 저장하여 AB_sum[key] = val1, CD_sum[-key] = val2라면 바로 val1*val2값을 sol에 더해주면 됩니다.

 

방법론 자체는 틀리지 않았다고 생각하는데, AB_sum.keys()에 존재하는 key의 부호가 반대인 값이 CD_sum.keys()에 있는지 확인하는 과정에서 시간 초과가 뜬 것 같습니다. dictionary의 key가 있는지 확인하는게 $O(1)$인 줄 알았는데 아닌가봅니다. 실제로 본 문제를 파이썬으로 푼 사람 중 keys in dict:를 활용해 풀었다고 하는 사람들이 있는데, 그 분들 중 한 분 코드를 그대로 돌려보니 시간 초과가 뜨더라구요. 1달 전 쯤 도합 3만개에 가까운 결과가 재채점 되었는데, 아마 그 때 재채점 통과를 실패하신 것 같습니다.