[백준] 3151 합이 0

🔗 문제 링크

https://www.acmicpc.net/problem/3151

💻 코드

제출 코드

어제 이어서 다시 풀어봤지만, 또 틀렸다. 😥😥
반례를 잡아내지 못한 것이 문제

import sys
input = sys.stdin.readline 

n = int(input())
nums = list(map(int,input().split()))
nums.sort()
result = 0

for i in range(n-2):
    start = nums[i]
    left,right = i+1,n-1
    if start > 0:
        break
        
    while left < right:
        tot = nums[left] + nums[right]
        left_tmp,right_tmp = 1,1
        if start + tot == 0: 
            
            while left+1 < right and nums[left] == nums[left+1]  :
                left_tmp += 1
                left += 1

            while left < right-1 and nums[right] == nums[right-1]:
                right_tmp += 1
                right -=1 

            result += left_tmp*right_tmp
            left += 1
            right -= 1

        elif start + tot < 0: # 음수면 left += 1
            left += 1
        else:
            right -= 1


print(result)

정답 코드

import sys
input = sys.stdin.readline 

n = int(input())
nums = list(map(int,input().split()))
nums.sort()
result = 0


for i in range(n-2):
    start = nums[i]
    left,right = i+1,n-1


    if start > 0:
        break
        
    while left < right:
        tot = nums[left] + nums[right]

        if start + tot == 0:
            
            if nums[left] == nums[right]:
                count = right - left + 1
                result += count * (count - 1) // 2
                break
            
            left_tmp,right_tmp = 1,1
            
            while left+1 < right and nums[left] == nums[left+1]  :
                left_tmp += 1
                left += 1

            while left < right-1 and nums[right] == nums[right-1]:
                right_tmp += 1
                right -=1 

            result += left_tmp*right_tmp
            left += 1
            right -= 1
            
        elif start + tot < 0: # 음수면 left += 1
            left += 1
        else:
            right -= 1



print(result)

6
-10 5 5 5 5 5 5 5

이 반례를 생각하지 못했다..
이럴 경우엔 조합으로 예외처리를 해줘야한다고 생각했다.
그런데 이렇게 코드를 복잡하게 구현할 수 밖에 없을까?

🤷‍♂ 풀이

나는 투포인터로 접근해서, start를 고정시키고 left,right를 이동하는 식으로 풀이했다.
대신 nums[left] == nums[right]라면 답이되는 경우의 수에서 값이 모두 동일한 것이므로 그 때는 nC2의 조합으로 개수를 구하게 하였다.
하지만 .. 내 풀이는 이해하기에 너무 복잡하여 다른 분의 풀이를 참고해봤다.

참고 풀이

# 합이 0
from bisect import bisect_left

N = int(input())
arr = list(map(int, input().split()))
arr.sort()

answer = 0
for i in range(len(arr)-2):
    left, right = i+1, N-1
    while left < right:
        result = arr[i]+arr[left]+arr[right]
        if result > 0:
            right -= 1
        else:
            if result == 0:
                if arr[left] == arr[right]:
                    answer += right - left
                else:
                    idx = bisect_left(arr, arr[right])
                    answer += right-idx+1
            left += 1

print(answer)

입력을 받고, 정렬을 한다.
start 변수는 필요가 없다. for i in range(N-2)의 i를 활용하면 되기 때문이다.
3명을 고르기 때문에 0부터 N-3까지 진행한다.
- left, right는 시작의 다음, 배열의 끝 원소로 초기화한다.
- result: 세 수의 합
- 값이 양수라면 값을 줄여야하므로 right -= 1
- 값이 0이거나 음수라면 값을 키워야하므로 left += 1
result == 0인 경우
- arr[left] == arr[right]라면, left와 right 사이의 모든 수가 같다는 것이다.
  - 이 때 조합의 수는 right - left가 된다.
  - 예를 들어 [1,1,1,1]이라면 left=1, right=3, 조합 3개
- arr[left] != arr[right]라면, bisect_left를 사용해 arr[right]의 첫 번째 위치를 찾는다.
  - 이 때 조합의 수는 right - idx + 1이다.
  - 예를 들어, [-8,3,3,5,5,5]라면, bisect_left(arr,5) = 3, 조합의 수는 5-3+1 = 3개
  - 세번째 수 nums[right]을 5로 고정한 조합의 개수가 3개라는 뜻이다.

right 포인터가 배열의 뒤쪽에서 시작하기에 같은 값을 가진 요소들이 연속적으로 나타나는 구간을 빠르게 처리한다.
left 포인터는 이동하면서 중복처리가 자동으로 해결되니까 right를 별도로 감소시키는 추가 로직이 필요하지 않다.
어렵다... 😫😫