[백준/Python] 3151번: 합이 0

문제

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

3151번: 합이 0

---

문제 요약

N명의 학생들이 각 -10000 ~ 10000까지의 코딩 실력을 갖고 있다.

 

세 명의 학생들을 골라 코딩 실력이 0이 되는 3인조를 만들 수 있는 경우의 수를 출력한다.

 

+ 코딩 실력이 같은 학생들이 있을 수 있다.

---

코드

N = int(input())

arr = sorted(list(map(int, input().split())))

cnt = 0

# 한 개의 수 선택 후 투 포인터를 통해 나머지 두 개의 수 구하기
for i in range(N - 2):
    left = i + 1
    right = N - 1
    max_idx = N
    # max_id는 세 수 중 가장 큰 수. 중복된 수라면 가장 왼쪽에 있는 수의 인덱스 (초기 값: N)
    while left < right:
        tmp = arr[left] + arr[right]

        if tmp < -arr[i]: # 합이 0보다 작은 경우
            left += 1

        elif tmp == -arr[i]: # 합이 0인 경우
            if arr[left] == arr[right]: 
            # left원소와 right원소가 같으면 사이에 있는 원소도 같은 값 (정렬되어 있기 때문)
                cnt += right - left
            else:
            # left원소와 right원소가 다른 경우
                if max_idx > right: # max_id가 제대로 설정되지 않은 경우만
                    max_idx = right
                    while max_idx >= 0 and arr[max_idx - 1] == arr[right]:
                        max_idx -= 1
                cnt += right - max_idx + 1
                
            left += 1

        else: # 합이 0보다 큰 경우
            right -= 1

print(cnt)

---

코드 설명

입력 받은 코딩 실력 배열을 오름차순 정렬한다. (투 포인터 방식을 사용하기 위해)

 

1. 0 ~ N - 3번째 학생까지 한 명을 고정으로 선택하고, 나머지 두 명을 투 포인터 방식으로 선택한다.

+ 첫 번째 학생의 코딩 실력: arr[i]+ 두 번째 학생의 코딩 실력: arr[left]+ 세 번째 학생의 코딩 실력: arr[right]

+ max_idx는 세 번째 학생과 같은 실력을 가진 학생들 중 가장 왼쪽에 있는 학생의 인덱스 (초기: N)

 

2.  세 학생의 실력 합이 0보다 작으면 left인덱스를 1 증가시키고, 크면 right인덱스를 1 감소시키며 탐색한다.

 

3. 세 학생의 실력 합이 0인 경우,

 

3 - 1. left원소와 right원소가 같으면 그 사이에 있는 값들도 같은 값이므로 right - left 개수만큼 cnt를 증가시킨다. 

ex) arr = [-6 3 3 3]

+ 정렬되어 있기 때문에 같은 값들이 모여 있다.

 

3 - 2. left원소와 right원소가 다를 경우, max_idx를 탐색하고 right - max_idx + 1 개수 만큼 cnt를 증가시킨다.

ex) arr = [-6 0 6 6 6]

 

3 - 3. left 인덱스를 1 증가시키고 다음 탐색을 진행한다.

 

4. cnt를 출력한다.