[Algorithm] 정렬된 배열에서 특정 수의 개수 구하기

문제

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있다.

이 때, 이 수열에서 x가 등장하는 횟수를 계산하시오.

 

풀이

입력되는 수열이 오름차순이고, 완전 탐색으로 몇개인지 셀 수 있지만,

시간복잡도에서 O(N)의 시간 복잡도를 가지게 된다.

 

이보다 이진검색트리를 사용하면 O(logN)의 시간 복잡도로 해결할 수 있다.

먼저, x보다 큰 수의 인덱스를 찾는다. 그리고 x보다 크거나 같은 수의 인덱스를 찾아 빼면 갯수를 확인할 수 있다.

 

왼쪽 인덱스는 타겟과 크거나 같은 수의 인덱스를 찾는다.

    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;
    }

 

전체소스코드

import java.util.*;

public class Main {

    public static int n;
    public static int x;
    public static int[] arr = new int[1000001];

    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;
    }

    public static void main(String[] args) {

        Scanner sc = new Scanner(System.in);
        n = sc.nextInt();
        x = sc.nextInt();
        for (int i = 0; i < n; i++) {
            arr[i] = sc.nextInt();
        }

        int left = leftIndex(0, n);
        int right = rightIndex(0, n);
        
        int result = right - left;

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