[백준 3151번] : 합이 0 파이썬 풀이

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

합이 0 

4 초

128 MB

6363

1494

1006

22.336%

문제

Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다. 코딩 실력이 좋으면 팀워크가 떨어지고, 팀워크가 좋을수록 코딩 실력이 떨어진다. 그리고 출전하고자 하는 대회는 코딩 실력과 팀워크 모두가 중요하다.

Elly는 그녀가 가르칠 수 있는 모든 학생들의 코딩 실력을 알고 있다. 각각의 코딩 실력 Ai는 -10000부터 10000 사이의 정수로 표시되어 있다. 그녀는 팀워크와 코딩 실력이 모두 적절한 팀을 만들기 위해, 세 팀원의 코딩 실력의 합이 0이 되는 팀을 만들고자 한다. 이러한 조건 하에, 그녀가 대회에 출전할 수 있는 팀을 얼마나 많이 만들 수 있는지를 계산하여라.

N명의 학생들의 코딩 실력 Ai가 -10000부터 10000사이의 정수로 주어질 때, 합이 0이 되는 3인조를 만들 수 있는 경우의 수를 구하여라.

입력

입력은 표준 입력으로 주어진다.

첫 번째 줄에 학생의 수 N이 입력된다. 두 번째 줄에 N개의 그녀가 가르칠 학생들의 코딩 실력인 Ai가 주어진다.

출력

표준 출력으로 그녀가 고를 수 있는 팀의 수를 하나의 정수로 출력한다.

제한

• 1 ≤ N ≤ 10000
• -10000 ≤ Ai ≤ 10000

예제 입력 1 복사

10
2 -5 2 3 -4 7 -4 0 1 -6

예제 출력 1 복사

6

힌트

예시에서 가능한 참가자 그룹은 아래와 같다.

(2, -5, 3), (2, 2, -4), (2, 2, -4), (-5, 2, 3), (3, -4, 1), (3, -4, 1)

두 개의 -4는 서로 다른 참가자를 나타내는 것에 유의하라. (2, 2, -4)와 (3, -4, 1)이 두 번씩 나타난다.

풀이

투 포인터와 이분 탐색을 활용하여 푸는 문제이다.

 

문제를 풀기 위한 기본적인 아이디어는 아래와 같다.

 

0. 먼저 배열을 오름차순으로 정렬한다.

 

1. 세개의 숫자 중 기준이 되는 한 숫자를 왼쪽에서 부터 결정한다. 편의상 이 수를 기준수라고 부르겠다.

 

2. 투 포인터를 활용하여 기준수의 오른쪽 수들중에서 숫자 2개를 결정한다.

두 수와 기준수를 더해서 0이 되어야 하므로, 두수의 합이 -기준수와 같아야 한다.

예를 들어 기준수가 -3이라면 나머지 두 수의 합은 3이어야 한다. 이 두 수를 결정하는 과정에서 투 포인터를 활용하며 예시를 들자면 아래와 같다.

 

-6 -5 -4 -4 0 1 2 2 3 7
select = -6; target = 6

start,end = 1,9 (배열의 index)
-5 + 7 = 2 < target ; start++
-4 + 7 = 3 < target ; start++
-4 + 7 = 3 < target ; start++
0 + 7 = 7 > target ; end--
0 + 3 = 3 < target ; start++
1 + 3 = 4 < target ; start++
2 + 3 = 5 < target ; start++
2 + 3 = 5 < target ; start++
start == end ; break

 

투 포인터 start와 end를 두고, 기준수의 오른쪽 부분 배열에서 왼쪽 끝과 오른쪽 끝을 각각 가리키게 한다.

start와 end가 가리키는 두 수를 더한 값이

target과 비교하여 더 작은 수라면

더 많이 더해야 한다는 의미이므로, start를 오른쪽으로 한칸 당긴다. 

target과 비교하여 더 큰 수라면

합을 조금 줄여야 한다는 의미이므로, end를 왼쪽으로 한칸 당긴다.

target과 비교하여 같은 수라면

같은 조건을 만족하는 수가 더 많이 있을 수 있으므로 추가 검증을 진행한다.

 

이 추가 검증 과정에서 추가적인 아이디어가 필요한데, 다음 예시를 통해 설명하겠다.

 

아래와 같은 배열이 있고, 기준수가 -8인 상황을 가정해보자.

 

start가 가리키는 수는 3으로 end는 5라면 data[start] + data[end] = 8로 우리가 찾고자 하는 수와 동일해진다.

배열을 살펴보면 5를 값으로 갖는 경우가 많이 등장하는데, 만약 start이후의 수가 모두 5인 경우를 상상해보면 end를 계속 -1해줘야 하므로 시간 복잡도는 결국 O(n)과 거의 동일해진다.

따라서 이러한 상황을 방지 하기 위해 bisect_left() 라이브러리를 활용하여 여러개의 5들 중 가장 왼쪽 인덱스에 해당하는 값을 알 수 있다면 end에서 그러한 인덱스를 빼주는 것으로 조건을 만족하는 전체 end의 수를 알 수 있다.

 

이번엔 아래와 같은 경우를 살펴보자.

기준수가 -10이고, start와 end는 각각 양 끝 값을 가리키고 있다.

 

배열은 정렬되어 있는 상황이므로, data[start]와 data[end]의 값이 같다면 추가적인 검증이 필요없이 start~end 사이의 값들도 모두 5라는 것을 알 수 있으므로, start가 1번째 인덱스인 상황에서 조건을 만족하는 end의 개수는 end-start개임을 알 수 있다.

이러한 규칙을 적용시키면서 start가 1에서 2,3,4, 등으로 옮겨가도록 하면 end포인터를 처음까지 당기는 연산을 수행 할 필요 없이 비교적 짧은 시간복잡도로 문제를 해결 할 수 있다.

 

이러한 아이디어를 코드로 옮기면 아래와 같다.

 

# https://www.acmicpc.net/problem/3151
'''
먼저 전체 데이터를 정렬
왼쪽에서 부터 수를 하나 지정, 해당 수와 더해서 0이 되도록하는 조건수 지정
해당 수의 오른쪽 수에서 부터 투포인터를 사용하여 투포인터를 통해 조건 수를 만들 수 있는지 확인
'''
import sys
from bisect import bisect_left
input = sys.stdin.readline
N = int(input())
data = list(map(int,input().split()))
data.sort()
result = 0 # 경우의 수(정답에 해당)
# 조건수 지정
for i in range(N-2):
    # i번째 수와 더해서 0이 될 수 있는 조건 수 지정
    target = -data[i]
    start,end = i+1,N-1 # 이분탐색의 양 끝점 지점(투 포인터 지정)
    while start<end:
        temp = data[start] + data[end]
        # 양 끝점을 더한 수가
        if temp == target:
            # 값이 같은 경우
            if data[start] == data[end]:
                result += (end-start)
            else: # 값이 다를 경우
                # 오른쪽에 해당하는 값의 개수만큼
                index = bisect_left(data,data[end]) # data[end]중 가장 왼쪽 인덱스 리턴
                result += end-index+1
            start+=1
        elif temp < target: 
            start+=1
        else: # temp > target
            end-=1

print(result)