[백준] 3151번: 합이 0

initial

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

import sys
input = sys.stdin.readline

n = int(input())
lst = list(map(int, input().split()))
lst.sort()

ans = 0

def search(left, right, cur_index):
    global ans
    while left < right:
        if left == cur_index:
            left += 1
            continue
        elif right == cur_index:
            right -= 1
            continue

        if lst[left] + lst[right] + lst[cur_index] == 0:
            ans += 1
            break
        elif lst[left] + lst[right] + lst[cur_index] < 0:
            left += 1
        else:
            right -= 1

for i in range(len(lst) - 2):
    search(0, len(lst) - 1, i)

print(ans)

while this passed the example, it coudlnt pass edge cases like

4
-2 -2 -2 4

정답 : 3
but mine was like 4

solution

The tricky part is to even though there are duplicate values, we still need to make a unique combination out of them. Look at the question example.

The key point is that my previous approach was making left and right pointers overlap with my current pointer at some point, which I needed to add continue conditions when that happened. But this below approach fixated current pointer while moving left and right pointers closer together.

When 3sum is 0 and lst[left]==lst[right] like [-4, 2, 2, 2], cur_index is 0, left is 1 and right is 3, we know that (1,3) and (2,3) are combinations with the cur_index of 0. So we can just add the length of right-left which is 2. You can think of it as like moving left pointer to the right till right pointer (1 to 2).

When 3sum is 0 and lst[left] is not equal to right like [-4, 1, 3, 3], and cur_index is 0, left is 1 and right is 3, we wanan find the leftmost index in this list that has value 3. That is index 2. So we know index 3 - index 2 + 1 = 2 values of 3 that satisfy that 3 sum is 0. So we add 2 to our answer. You can think of it as like moving right pointer to the left till its value to the left is not equal to its own value.

2sum is easy but 3 sum is hard.

from bisect import bisect_left
import sys

input =sys.stdin.readline

n = int(input())
lst = list(map(int, input().split()))
lst.sort()

ans = 0

def search(left, right, cur_index):
    global ans
    while left < right:
        if lst[left] + lst[right] + lst[cur_index] > 0:
            right -= 1
        else:
            if lst[left] + lst[right] + lst[cur_index] == 0:
                if lst[left] == lst[right]:
                    ans += right - left
                else:
                    leftmost_idx_of_right_val = bisect_left(lst, lst[right])
                    ans += right - leftmost_idx_of_right_val + 1
            left += 1

for i in range(len(lst) - 2):
    search(i + 1, len(lst) - 1, i)

print(ans)

complexity

Bisect_left is log n but it is in a nested loop so it is n^2 log n. Space is linear