[백준] 3151번 - 합이 0 (java)

Baekjoon 3151 - 합이 0 (클릭 시 이동)

문제

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

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

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

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
4 초	128 MB	1635	369	223	21.401%

입력

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

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

출력

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

풀이

최대 경우의 수가 존재할 때는 N 이 10000이고 배열이 모두 0으로 주어진 경우이다.

이 때 경우의 수는 10000C3 = 166,616,670,000이므로 long타입으로 정답을 출력해야 한다.

1. 입력값을 오름차순으로 정렬한다.
2. 가장 작은 값(0번 인덱스)을 fix 시키고 1번인덱스와 마지막 인덱스를 start, end로 놓는다.
3. 배열의 0, start, end의 값이 0인지 확인한다
4. 합이 0보다 크면 end를 줄인다, 0보다 작거나 같으면 start를 늘린다.

효율 고려

같은 값이 존재할 수 있다. 만약 [-8, -2, -2, -2, 10, 10, 10, 10 ] 이라는 배열이 존재한다고 가정해 보자.

만들어 질 수 있는 부분집합은 {-8, -2, 10}밖에 존재하지 않지만 각각의 인덱스는 다른 경우가 된다.

따라서 -2의 갯수인 3과 10의 갯수인 4를 곱한 12가지의 경우가 존재한다.

또는 [-10, 5, 5, 5, 5, 5, 5, 5]라는 배열이 주어진다면 -10이 fix 되어 있을 때 start와 end가 가리키는 배열의 값은 같고 세개의 수의 합은 0이 된다. 이 때 경우의 수를 구하는 과정은 아래와 같다.

if(합 == 0){
    if(data[start] == data[end]){
        answer += comb(end-start-1);    // nC2
    }
}

static int comb(int n) {
        return n * (n - 1) / 2;
    }

최종 소스코드

import java.io.*;
import java.util.*;

public class Main {
    static int N;
    static long answer = 0L;
    static int[] data;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = null;

        N = Integer.parseInt(br.readLine());
        data = new int[N];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            data[i] = Integer.parseInt(st.nextToken());
        }
        Arrays.sort(data);
        for (int i = 0; i < N; i++) {
            if (data[i] > 0) break;
            int start = i + 1;
            int end = N - 1;

            while (start < end) {
                int s = 1;
                int e = 1;
                int current = data[i] + data[start] + data[end];
                if (current == 0) {
                    if (data[start] == data[end]) {    // start == end일 경우
                        answer += comb(end - start + 1);
                        break;
                    }
                    // start의 다음 값이 같은 경우
                    while (start + 1 < end && data[start] == data[start + 1]) {
                        s++;
                        start++;
                    }
                    // end의 다음 값이 같은 경우
                    while (start < end - 1 && data[end] == data[end - 1]) {
                        e++;
                        end--;
                    }

                    answer += s * e;
                }

                if (current > 0)
                    end--;
                else start++;
            }
        }
        System.out.println(answer);
    }

    static int comb(int n) {
        return n * (n - 1) / 2;
    }
}