뒤죽박죽 코딩 이야기

2021년 9월 04일 토요일 - 으윽.. 개강... 개강해....

오늘 올려볼 문제는 3151번 합이0 이라는 문제이다.

사진을 클릭하면 해당 문제로 이동합니다.

오늘은 이렇게 오래 걸릴지 몰랐던 쉬워보이는 문제를 가져왔다.

자칫하면 놓치기 쉬운 개념들이 매우 많았다

입력 예제

문제가 어려운 것에 비해서 입력은 간단하다...!

풀이 및 코드

이번 풀이는 텍스트로만은 설명하기 힘들것 같아서 이미지를 준비해왔다! 천천히 살펴보자.

그 전에 문제에 중요한 조건을 하나 보자.

힌트

입력으로 들어온 수는 같은 수더라도 다른 사람으로 취급한다는 것이다. 이 점 유의하자.

먼저 들어온 수들을 정렬해서 음수 | 0 | 양수 와 같은 형태로 생각해보자.

풀이 1

그럼 이런 범위 안에서 첫번째 수와 두번째 수를 선택할 수 있을 것이다.

첫번째 수는 음수에서만 선택하고고 두번째 수는 첫번째 수보다 1 큰 수에서 선택한다. (여기서 수는 인덱스)

풀이 2

이렇게 범위를 나누는 이유는 앞에 두 수를 찾은 뒤에 세번째 수를 이진 탐색으로 찾으면 O(n^2)의 시간복잡도로 문제를 풀 수 있기 때문이다.

여기까지는 쉽게 이해될 것이라고 생각한다.

풀이 3

자 이제 첫번째수를 모두 선택했다고 가정해보자. 그럼 모든 경우의 수를 찾은 것일까?

아니다. 0으로만 이루어진 팀에 대해서 예외처리를 해주어야한다.

예외처리는 그냥 조합 공식만 쓰면 되므로 이는 쉽다. 물론 0이 3개 이상일 때만 예외처리를 해주면 된다.

풀이 4

아까 세번째 수를 고르고 나서 결과값을 구할 때 주의해야할 점이 하나 있다.

세번째 수가 여러개 존재한다면 이에 대해서 모두 더해주어야 하기 때문이다.

그런데 이를 선형적으로 루프문을 돌려서 모두 구하게 되면 O(n^3)의 시간복잡도가 되므로 시간초과가 나게 될 것이다.

그래서 해주어야 할 것이 수의 중복 개수를 저장하고, 중복되는 수의 처음 나오는 인덱스를 저장해야 한다.

그렇다면 위, 아래 두가지 상황에 대해서 간단한 수학 공식으로 결과값을 구할 수 있게 된다.

풀이 5

이제는 코드를 봐보자! (이번 코드는 좀 많이 더럽다... 가독성이 떨어지니 주석을 잘 참고하길 바란다.)

풀이코드

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.Arrays;

public class Main {
    // n 저장(글로벌로 씀)
    public static int n;
    // 입력 저장할 배열, 중복되는 수의 가장 처음 나오는 index 저장하는 배열, 수의 중복 개수를 저장할 배열
    public static int[] arr, location, cnt;

    public static void main(String[] args) throws NumberFormatException, IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        // n 입력 받음
        n = Integer.parseInt(br.readLine());

        // 저장할 배열을 n + 1의 크기를 할당(처리 편하게 할려고 크게 잡는 것임)
        arr = new int[n + 1];

        // 위치와 중복을 저장할 배열을 200005의 크기를 할당 (배열에는 음수의 인덱스가 없으므로 10000을 0이라고 치고 계산하기 위해서 그런 것, 이것 또한 처리하기 편하게 크게 잡음)
        location = new int[20005];
        cnt = new int[20005];

        // 0 개수를 저장함
        int zero = 0;

        // 입력 받음
        String[] input = br.readLine().split(" ");

        // int형으로 캐스팅하면서 0의 개수 및 수의 중복 개수도 저장
        for(int i = 0; i < n; i++)
        {
            arr[i] = Integer.parseInt(input[i]);
            if(arr[i] == 0)
                zero++;
            cnt[10000 + arr[i]]++;
        }
        // 배열을 1크게 잡았으므로 한 칸이 남게 되는데 기본적으로 0이 저장되어 있다. 이 상태로 정렬을 하게 되면
        // 원하지 않는 0이 섞여서 저장되게 되므로 이 값을 입력으로 받을 수 없는 큰 값으로 만들어서
        // 맨 오른쪽에 위치하게끔 해주는 작업이다.
        arr[n] = 20000;

        // 오름차순으로 정렬을 해준다.
        Arrays.sort(arr);


        // 중복되는 수가 처음 나오는 인덱스를 저장하기 위한 작업
        int pre = arr[0];
        location[10000 + arr[0]] = 0;
        for(int i = 1; i < n; i++)
        {
            if(pre != arr[i])
            {
                pre = arr[i];
                location[10000 + arr[i]] = i;
            }
        }

        // 결과를 저장할 변수
        long result = 0;

        // 첫번째로 선택할 수는 arr[i], 음수에서만 선택한다.
        for(int i = 0; i < n && arr[i] < 0; i++)
        {
            // 두번째로 선택할 수는 arr[j], i + 1 부터 선택한다.
            for(int j = i + 1; j < n; j++)
            {
                // 두 값의 합이 음수라면
                if(arr[i] + arr[j] < 0)
                {
                    // 두 값의 합과 더했을 때 0이 되는 수를 구해준다.
                    int num = -1 * (arr[i] + arr[j]);
                    // 이진탐색으로 해당하는 수가 인덱스 j+1부터 배열의 끝까지 존재하는 확인한다.
                    if(search(num, j + 1))
                    {
                        // 존재한다면 그 수가 처음 나오는 인덱스가 3번째로 선택할 수 있는 인덱스보다 작은지 확인한다
                        // 작다면
                        if(location[10000 + num] < j + 1)
                        {
                            // (수의 중복 개수 + 수의 인덱스 위치 - 3번째로 선택할 수 있는 인덱스)를 결과값에 더해준다.
                            result += cnt[10000 + num] + location[10000 + num] - j - 1;
                        }
                        // 크다면
                        else
                        {
                            // 결과값에 수의 중복 개수를 더해준다.
                            result += cnt[10000 + num];
                        }
                    }
                }
                else
                {
                    break;
                }
            }
        }

        // 0이 3개 이상 나왔다면
        if(zero >= 3)
            // 조합 공식을 사용하여 결과값에 더해준다.
            result += combination(zero);

        bw.write(result + "\n");
        bw.flush();
    }

    // 조합 함수
    public static long combination(int zero)
    {
        // (0의 개수 C 3)의 값을 리턴한다.
        return zero * --zero * --zero / 1 / 2 / 3;
    }

    // 이진 탐색, 찾을 값과 탐색해야하는 인덱스를 파라매터로 받는다.
    public static boolean search(int num, int index)
    {
        // index부터 탐색을 한다.
        int left = index, right = n - 1, middle;

        while(left <= right)
        {
            middle = (left + right) / 2;

            if(arr[middle] == num)
            {
                return true;
            }
            else if(arr[middle] > num)
            {
                right = middle - 1;
            }
            else
            {
                left = middle + 1;
            }
        }

        return false;
    }

}

제출 화면

쉬워보였는데 신경쓸게 너무 많아서 많이 틀리고 조금 오래 걸렸던 문제이다...

분명 나처럼 고려해야 할 부분들을 미처 생각하지 못한 사람들이 많을 것 같아서 블로그에 올려본다.