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

문제 설명

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있을 때, 이 수열에서 x가 등장하는 횟수를 계산하세요.

예를 들어 수열 {1, 1, 2, 2, 2, 2, 3}이 있을 때 x=2라면, 4를 출력

시간복잡도 O(log N)으로 알고리즘을 설계하지 않으면 '시간초과 판정'

 

입력조건

• 첫째 줄에 N과 x
• 둘째 줄에 n개의 원소

출력조건

• 수열에서 값이 x인 원소의 개수
• 값이 x인 원소가 없다면 -1

입출력 예시

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

 

 

 

 

문제 풀이

• 시간복잡도 O(log N) 로 설계해야하기 때문에 순차 탐색으로는 풀 수 없다. → 이진탐색으로 풀이
• x가 처음으로 등장하는 위치와 마지막 위치를 찾아서 빼주면 x의 개수를 구할 수 있다.
• firstIndex() : x가 처음으로 등장하는 위치를 찾는 이진 탐색 함수
• lastIndex() : x의 마지막 위치를 찾는 이진 탐색 함수

import java.io.*;
import java.util.StringTokenizer;

public class Main {

    public static int x, count = 0;
    public static int[] array;

    //x가 처음 위치를 찾는 이진 탐색
    public static int firstIndex(int start, int end) {
        while(start < end){
            int mid = (start + end) / 2;
            if(array[mid] >= x){
                end = mid;
            } else {
                start = mid + 1;
            }
        }
        return end;
    }

    //x의 마지막 위치를 찾는 이진 탐색
    public static int lastIndex(int start, int end) {
        while(start < end){
            int mid = (start + end) / 2;
            if(array[mid] > x){
                end = mid;
            } else {
                start = mid + 1;
            }
        }
        return end;
    }

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

        st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        x = Integer.parseInt(st.nextToken());

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

        int first = firstIndex(0, n);
        int last = lastIndex(0, n);

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

 

 

처음에는 아래와 같이 재귀 함수를 사용해서 이진 탐색을 구현했다.

• array[mid] == x이면 mid를 중심으로 왼쪽, 오른쪽 영역에 x가 있는지 재귀함수를 호출해서 탐색한다.
• 하지만 이방법은 시간복잡도가 O(log N)을 만족하지 못하는 것 같다...

public static int binary_search(int start, int end) {
    if (start > end) {
        return -1;
    }

    int mid = (start + end) / 2;
    if (array[mid] == x) {  //x와 같다면
        count++;
        binary_search(mid + 1, end);    //mid를 중심으로 오른쪽 영역에 x가 있는지 확인
        binary_search(start, mid - 1);  //mid를 중심으로 왼쪽 영역에 x가 있는지 확인
    } else if (array[mid] < x) {
        binary_search(mid + 1, end);
    } else {
        binary_search(start, mid - 1);
    }
    return count;
}

---

[참고서적]

이것이 취업을 위한 코딩 테스트다 with 파이썬 by 나동빈