[이진 탐색] 이코테 정렬된 배열에서 특정 수의 개수 구하기(Java)

📚 문제

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이때 이 수열에서 x가 등장하는 횟수를 계산하세요. 값이 x인 원소가 하나도 없다면 -1을 출력합니다.  (단, 이 문제는 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 '시간 초과' 판정을 받습니다.)

 

첫째 줄에 N (0<=N<=1,000,000) , x (-10^9 <= x <= 10^9)가 정수 형태로 공백으로 구분되어 입력됩니다.

둘째 줄에 N개의 원소가 정수 형태로 공백으로 구분되어 입력됩니다.

 

입력 예시	출력 예시
7 2 1 1 2 2 2 2 3	4
7 4 1 1 2 2 2 2 3	-1

 

📝 문제 해결

모든 원소가 정렬이 상태로 입력되므로, 수열 내에 x가 존재한다면 연속적으로 나열되어 있을 것으로 예상할 수 있다. 따라서 x가 처음 등장하는 인덱스와 x가 마지막으로 등장하는 인덱스를 각각 계산한 뒤에, 그 인덱스의 차이를 계산하면 문제를 해결할 수 있다. 그러므로 이진 탐색 함수를 2개 작성하여 문제를 해결한다. (이러한 방법은 이진 탐색을 요구하는 고난이도 문제에서 자주 사용할 수 있는 테크닉이다.)

 

하나는 데이터가 존재한다면 가장 첫 번째 위치를 찾는 이진 탐색 함수이며,

다른 하나는 데이터가 존재한다면 가장 마지막 위치를 찾는 이진 탐색 함수이다.

 

• 첫 번째 위치를 찾는 이진 탐색 함수 = x보다 크거나 같은 인덱스를 찾는다. = 중간값이 x보다 크거나 같으면 끝값을 중간값으로 이동 
• 마지막 위치를 찾는 이진 탐색 함수 = x보다 큰 인덱스를 찾는다 = 중간값이 x보다 큰 경우 끝값을 중간으로 이동 

 

💻 코드

package search;

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

public class Main {
    public static int N, x;
    public static int[] arr;

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

        N = Integer.parseInt(st.nextToken());
        x = Integer.parseInt(st.nextToken());
        arr = new int[N];

        st = new StringTokenizer(br.readLine());
        for(int i=0; i<N; i++){
            arr[i] = Integer.parseInt(st.nextToken());
        }

        int left = leftIndex(0, N);
        int right = rightIndex(0, N);

        System.out.println(right-left == 0 ? -1 : right-left);
    }

    public static int leftIndex(int start, int end){
        while(start < end){
            int mid = (start+end)/2;
            if (arr[mid] >= x) end = mid;
            else start = mid+1;
        }
        return end;
    }

    public static int rightIndex(int start, int end){
        while (start < end){
            int mid = (start+end)/2;
            if(arr[mid] > x) end = mid;
            else start = mid+1;
        }
        return end;
    }
}