[python] 백준 3151 합이 0

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

3151번: 합이 0

 

 

 

문제 해결

• 3개의 수를 고르는 문제는 전형적으로 A[i] + A[j] + A[k] = c 에서 A[j] + A[k] = c - A[i]로 바꾼다.
• c= 0이므로 A[j]+A[k] = -A[i]를 찾는 문제
• 그러면 A[i]를 고정시키고 j와 k를 찾는 것이 편하다.
• 처음에 j= i+1 k=n-1으로 정하고 A[j]+A[k] >A[i]면 k를 내리고 A[j] + A[k] < A[i]이면 j를 올리는 식으로 j,k를 찾으면 된다.

 

CODE

def solve(s,e,g):
    global answer
    max_idx = n
    while s<e:
        temp = A[s] + A[e]
        if temp < g:
            s += 1
        elif temp == g:
            if A[s] == A[e]:
                answer += e-s
            else:
                if max_idx > e:
                    max_idx = e
                    while max_idx > s and A[max_idx-1] == A[e]:
                        max_idx -= 1
                answer += e - max_idx + 1
            s += 1
        else:
            e -= 1
if __name__=='__main__':
    n = int(input())
    A = list(map(int, input().split()))
    A.sort()
    answer = 0
    for i in range(n-2):
        start = i+1
        end = n-1
        goal = -A[i]
        solve(start,end,goal)
    print(answer)