[Problem Solving] BOJ 3151 : 합이 0

문제 링크 : https://www.acmicpc.net/problem/3151

 

세 수의 합이 0이 되는 경우의 수를 구하는 문제이다. 수가 동일하지만 다른 위치에 있다면 다른 수로 취급한다.

 

우선 배열 A를 정렬한다.

left, mid, right라는 세 개의 지점을 잡고, A[left] + A[mid] + A[right]가 0되는 경우의 수를 모두 찾으면 된다.

세 수의 대소 조건은 0 <= left < mid < right <= N-1 이다.

 

left는 0부터 N-3까지 가능하다.

left를 먼저 잡고, mid = left + 1, right = N - 1 으로 left 이하 부분은 제외한 양쪽 끝으로 설정한다.

다음 과정을 mid < right일 동안 반복한다.

sum = A[left] + A[mid] + A[right]라고 하자.

1. sum < 0이라면 sum을 증가시키기 위해서 mid를 오른쪽 한칸 이동시킨다.

2. sum > 0이라면 sum을 감소시키이 위해서 right을 왼쪽으로 한칸 이동시킨다.

3. sum == 0이면서 A[mid] == A[right]라면 mid와 right사이의 수는 모두 같은 수가 된다. 하지만 자리가 다르면 다 다른 수 이기 때문에 mid와 right 사이에 있는 수 중에서 두 개를 고르는 경우의 수를 계산하면 된다. mid와 right 사이에 있는 수의 개수는 right - mid + 1이고, 그것들 중에 두 개를 고르는 경우의 수는 (right - mid + 1) * (right - mid) / 2 이다. 이를 answer에 더해준다. 그 다음 break해준다.

sum == 0이면서 A[mid] != A[right]라면, mid보다 큰 부분에서 A[mid]와 같은 수들이 있을 수도 있고, right보다 작은 부분에서 A[right]와 같은 수들이 있을 수 있다. 각각의 수를 구해서 곱하면 만들 수 있는 조합의 경우의 수를 구할 수 있다. 각각의 수를 midCnt와 rightCnt에 저장해서 곱한 midCnt * rightCnt를 answer에 더해준다. 그 다음 mid를 1증가, right를 1감소 시킨다.

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

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

    public static void input() throws IOException
    {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());
        A = new int[N];
        String[] line = br.readLine().split(" ");
        for(int i = 0; i < N; i++)
        {
            A[i] = Integer.parseInt(line[i]);
        }
        br.close();
    }

    public static void solution()
    {
        Arrays.sort(A);
        for(int left = 0; left < N - 2; left++)
        {
            if(A[left] > 0)
                break;

            int mid = left + 1;
            int right = N - 1;
            while(mid < right)
            {
                int sum = A[left] + A[mid] + A[right];
                if(sum < 0)
                {
                    ++mid;
                }
                else if(sum > 0)
                {
                    --right;
                }
                else if(sum == 0)
                {
                    if(A[mid] == A[right])
                    {
                        answer += (right - mid + 1) * (right - mid) / 2;
                        break;
                    }
                    else
                    {
                        int midCnt = 1;
                        int rightCnt = 1;
                        while(mid + 1 < right && A[mid + 1] == A[mid])
                        {
                            midCnt += 1;
                            ++mid;
                        }
                        while(right - 1 > mid && A[right - 1] == A[right])
                        {
                            rightCnt += 1;
                            --right;
                        }
                        answer += midCnt * rightCnt;
                        ++mid;
                        --right;
                    }
                }
            }
        }
    }

    public static void output()
    {
        System.out.println(answer);
    }

    public static void main(String[] args) throws IOException
    {
        input();
        solution();
        output();
    }
}