백준 3151 합이 0 Kotlin (투 포인터)

문제 출처 : https://www.acmicpc.net/problem/3151

3151번: 합이 0

문제

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

Elly는 그녀가 가르칠 수 있는 모든 학생들의 코딩 실력을 알고 있다. 각각의 코딩 실력 Ai는 -10000부터 10000 사이의 정수로 표시되어 있다. 그녀는 팀워크와 코딩 실력이 모두 적절한 팀을 만들기 위해, 세 팀원의 코딩 실력의 합이 0이 되는 팀을 만들고자

한다. 이러한 조건 하에, 그녀가 대회에 출전할 수 있는 팀을 얼마나 많이 만들 수 있는지를 계산하여라.

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

입력

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

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

출력

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

제한

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

힌트

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

(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)이 두 번씩 나타난다.

알고리즘 분류

• 브루트포스 알고리즘
• 정렬
• 이분 탐색
• 두 포인터

풀이

입력 값을 보아하니 투 포인터로 풀어야할 것 같은 느낌이 빡 온다.

하지만 3개의 수를 골라야 하므로, 단순히 주어진 입력을 정렬하고, 바로 투 포인터를 적용할 순 없다.

일반 조합 또한 10000C3으로 시간 내에 풀 수 없다.

 

그러면 어떻게 시간 내에 답을 찾을 수 있을까?

 

6 

-6 -6 1 2 5 5

가 있다고 하자.

입력 받은 값들을 정렬시킨 모양이다.

 

-6을 선택했다고 했을 때, 나머지 두 개의 수를 선택해야 한다.

-6을 선택했다면 나머지 두 개의 수는 6이 되어야 한다.

오오오!

-6을 선택했을 때, 나머지 -6 1 2 5 5 에서 두 개 값의 합이6인 경우를 찾는 데에 투 포인터를 적용하면 된다!!

 

즉 i = (0~n-3)를 선택하고, left를 i+1, right를 n-1로 두고, arr[left]+arr[right] + arr[i]이 0보다 작으면 left++,

0보다 크면 right--를 해서 답을 찾아나가면 된다.

 

arr[left]+arr[right] + arr[i] ==0일 때는, 다음 경우를 고려해야 한다.

1.

6

-6 -6 1 1 2 5

2.

6

-6 -6 1 2 5 5

1.의 경우는 다음 left가 같은 경우이고, 이 경우는 arr[i] = arr[left] + arr[right]일 때도 디폴트로 left를 증가시키는 것으로 해결 가능하다.

 

2.의 경우는 다음 right가 같은 경우를 세어주고 cnt에 저장한 후, answer에 1을 더하는 것이 아닌 cnt를 더해준다.

다음 right가 같은 경우를 매번 센다면 시간 초과가 뜨기 때문에, cnt를 이미 구했다면 재사용하자.  

 

 

코드

fun main() = with(System.out.bufferedWriter()){
    val br = System.`in`.bufferedReader()
    val n = br.readLine().toInt()
    //입력값 정렬
    val input = br.readLine().split(' ').map { it.toInt() }.sorted()
    var answer=0L
    //하나 고정, 2개 더한 값(투 포인터)
    for(i in 0 until input.size-2) {

        var left = i+1
        var right = input.size-1
        var cnt=0
        while(left<right ){

            //합이 0일 때
            if(input[i]+input[left]+input[right]==0){
                if(input[left]==input[right]){
                    answer+=right-left
                    cnt=0
                }
                else {
                //cnt 한 번 구한 건 재사용
                    if(cnt==0) {
                        var idx = right
                        while (idx > left && input[i] + input[left] + input[idx--] == 0) {
                            cnt++
                        }
                    }
                    answer+=cnt
                }
                left++
            }
            //합이 0보다 작을 때
            else if(input[i]+input[left]+input[right]<0){
                left++
                cnt=0
            }
            //합이 0보다 클 때
            else{
                right--
                cnt=0
            }
        }
    }

    write("$answer")

    close()
}