Softeer HSAT 7 (Java)

문제 설명

정수 배열에서 세 가지 원소를 골라 가운데값(중앙값)을 구하는 상황입니다. 문제의 목적은 세 가지 원소를 고를때 이러한 특정 중앙값을 구할 수 있는 경우의 수가 얼마나 되는지 구합니다.

배열의 크기는 n 이며
이러한 행위가 q 번 진행됩니다.

평균이 아니라 중앙값

제약조건

• 1 ≤ n ≤ 50,000
• 1 ≤ q ≤ 200,000
• 1 ≤ 자동차의 연비 ≤ 1,000,000,000
• 1 ≤ mi ≤ 1,000,000,000 (i는 1 이상 q 이하입니다. 즉, mi 는 각 질의에 대응하는 중앙값을 의미합니다.)

접근

세 개 중의 중앙값을 구하다보니 이 경우의 수는
원하는 값보다 작은 원소의 수 x 원하는 값보다 큰 원소의 수 가 되겠습니다.

쟁점은 최대 200,000의 q에서 최대 50,000개의 원소가 있는 배열에서 어떻게 원하는 정보를 빠르게 얻을 것일까 입니다.

기본적으론 O(n) 의 시간에 값을 구할 수 있습니다.
그렇게되면 200,000 * 50,000 으로 상당히 큰 값입니다.

조회성능을 올리기 위해 binary search 를 고려할 수 있습니다.

O(nlgn) 으로 정렬을 하고 binary search 를 통해 원하는 정보를 O(lgn) 에 구할 수 있습니다. 즉 이 문제를 O(q*lgn + nlgn) 에 해결할 수 있습니다.

binary search 를 통해 어떻게 값을 구할 수 있느냐도 살펴봐야합니다.

그 중에서도 중복된 값이 있을 때 조금 귀찮아질 수 있습니다.

파이썬에선 bisect_left, right 같은 메소드로 중복된 구간을 다루기가 편하지만 자바에선 직접 구현해줍니다.

Arrays.binarySearch

    /**
     * Searches a range of
     * the specified array of ints for the specified value using the
     * binary search algorithm.
     * The range must be sorted (as
     * by the {@link #sort(int[], int, int)} method)
     * prior to making this call.  If it
     * is not sorted, the results are undefined.  If the range contains
     * multiple elements with the specified value, there is no guarantee which
     * one will be found.
     *
     * @param a the array to be searched
     * @param fromIndex the index of the first element (inclusive) to be
     *          searched
     * @param toIndex the index of the last element (exclusive) to be searched
     * @param key the value to be searched for
     * @return index of the search key, if it is contained in the array
     *         within the specified range;
     *         otherwise, <code>(-(<i>insertion point</i>) - 1)</code>.  The
     *         <i>insertion point</i> is defined as the point at which the
     *         key would be inserted into the array: the index of the first
     *         element in the range greater than the key,
     *         or {@code toIndex} if all
     *         elements in the range are less than the specified key.  Note
     *         that this guarantees that the return value will be &gt;= 0 if
     *         and only if the key is found.
     * @throws IllegalArgumentException
     *         if {@code fromIndex > toIndex}
     * @throws ArrayIndexOutOfBoundsException
     *         if {@code fromIndex < 0 or toIndex > a.length}
     * @since 1.6
     */
    public static int binarySearch(int[] a, int fromIndex, int toIndex,
                                   int key) {
        rangeCheck(a.length, fromIndex, toIndex);
        return binarySearch0(a, fromIndex, toIndex, key);
    }

1. 내용을 보자면 원하는 값이 없다면 음수값을 내려줍니다.
   원하는 값이 없다면 중앙값으로 둘 수 없으니 무조건 0입니다.

2. 중복된 값이 있다면 중복된 값 중 임의의 인덱스를 반환합니다.

구현 코드

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

public class SofteerHsat7 {
    public static void main(String[] args) throws IOException {
        BufferedReader bR = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bW = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer sT = new StringTokenizer(bR.readLine());
        int n = Integer.parseInt(sT.nextToken());
        int q = Integer.parseInt(sT.nextToken());
        
        sT = new StringTokenizer(bR.readLine());
        int[] nums = new int[n];

        for(int i = 0; i < n; i++) {
            nums[i] = Integer.parseInt(sT.nextToken());
        }

        Arrays.sort(nums);

        for(int i = 0; i < q; i++) {
            int _q = Integer.parseInt(bR.readLine());
            int left = Arrays.binarySearch(nums, _q);

            if (left < 0) {
                bW.write(0 + "\n");
                continue;
            }

            int right = nums.length - binarySearchRight(nums, _q) - 1;
            bW.write(left * right + "\n");
        }

        bW.flush();
        bR.close();
        bW.close();
    }

    private static int binarySearchLeft(int[] nums, int key) {
        int idx = Arrays.binarySearch(nums, key);

        if (idx < 0) {
            return idx;
        } else {
            while (idx > 0 && nums[idx - 1] == key) {
                idx--;
            }
            return idx;
        }
    }

    private static int binarySearchRight(int[] nums, int key) {
        int idx = Arrays.binarySearch(nums, key);

        if (idx < 0) {
            return idx;
        } else {
            while (idx < nums.length - 1 && nums[idx + 1] == key) {
                idx++;
            }
            return idx;
        }
    }
}