[99클럽 코테 스터디 3일차 TIL / 백준] 11663- 선분 위의 점 (JAVA)

99클럽 코테스터디 3일차 TIL
KeyWord : BinarySearch, upper bound, lower bound

문제

일차원 좌표상의 점 N개와 선분 M개가 주어진다. 이때, 각각의 선분 위에 입력으로 주어진 점이 몇 개 있는지 구하는 프로그램을 작성하시오.

입력

첫째 줄에 점의 개수 N과 선분의 개수 M이 주어진다. (1 ≤ N, M ≤ 100,000) 둘째 줄에는 점의 좌표가 주어진다. 두 점이 같은 좌표를 가지는 경우는 없다. 셋째 줄부터 M개의 줄에는 선분의 시작점과 끝점이 주어진다. 입력으로 주어지는 모든 좌표는 1,000,000,000보다 작거나 같은 자연수이다.

출력

입력으로 주어진 각각의 선분 마다, 선분 위에 입력으로 주어진 점이 몇 개 있는지 출력한다.

 

---

 

풀이

이분 탐색을 이용하여 해결할 수 있는 문제다.

1. 점의 좌표를 정렬하여 배열에 저장한다.
2. 선분의 시작점에서 가장 가까운 점, 끝점의 위치에서 가장 가까운 점의 인덱스를 구한다.
3. 인덱스를 이용해 점의 개수를 구한다.

여기서 이분탐색이 필요한 부분은 2번이다.

하지만 여기서 문제는 선분의 시작점 혹은 끝점이 좌표 배열에 속하지 않을 수 있다는 점이다. 즉, 우리가 배열에서 탐색해야 하는 것은 시작점 혹은 끝점이 아니다.

예를 들어 dots = [1, 3, 10, 20, 30]에서 선분이 2 15로 주어졌다면, 우리가 찾아야 하는 점은 2와 가장 가까운 점과 15와 가장 가까운 점이 된다.

처음에는 혹시나 이렇게 해도 되려나? 싶어서 넣어봤으나 당연하게도 시간초과가 나타났다.

for (int i = 0; i < m; i++) {
    stk = new StringTokenizer(br.readLine());

    int start = Integer.parseInt(stk.nextToken());
    int end = Integer.parseInt(stk.nextToken());

    int idx1;
    while((idx1=Arrays.binarySearch(dots, start)) <= -1) {
        start++;
    }

    int idx2;
    while((idx2=Arrays.binarySearch(dots, end)) <= -1)
        end--;

    sb.append(idx2 - idx1 + 1).append('\n');
}

 

dots[10000000, 10000001]이고, 선분이 1 10000000이라면, 1부터 10000000까지 전부 이진탐색을 돌려야 하니 시간 초과가 나타나는 건 당연하다.

이 문제를 해결하는 테크닉은 UpperBound와 LowerBound를 이용하는 것이다.

lowerBound는 찾고자 하는 key 이상의 값 중 최솟값을 찾는 것이고 upperBound는 key를 초과하는 값 중 최솟값을 찾는 방법이다.

dots = [1, 3, 10, 20, 30]에서 선분이 10을 찾는다고 하자. lowerBound라면 10을 찾을 것이고, upperBound라면 20을 찾을 것이다.

이 방법을 이용하면 된다. 시작점은 lowerBound, 끝점을 upperBound로 탐색하여 인덱스를 반환한다면 그 사이에 속한 점의 개수를 쉽게 알아낼 수 있다.

선분 2 15라면 idx1 = 1, idx2 = 3이 반환될 것이고, 이 사이에 속한 점의 개수는 idx2 - idx1이 된다.

 

코드

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

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

        int n, m;  
        n = Integer.parseInt(stk.nextToken());  
        m = Integer.parseInt(stk.nextToken());  

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

        Arrays.sort(dots);  

        for (int i = 0; i < m; i++) {  
            stk = new StringTokenizer(br.readLine());  

            int start = Integer.parseInt(stk.nextToken());  
            int end = Integer.parseInt(stk.nextToken());  

            int idx1 = lowerBound(dots, start);  
            int idx2 = upperBound(dots, end);  

            sb.append(idx2 - idx1).append('\n');  
        }  

        System.out.println(sb);  
    }  

    private static int lowerBound(int[] arr, int key) {  
        int lo = 0;  
        int hi = arr.length;  

        while (lo < hi) {  
            int mid = lo + (hi - lo) / 2;  

            if (key <= arr[mid]) {  
                hi = mid;  
            } else {  
                lo = mid + 1;  
            }  
        }  
        return lo;  
    }  

    private static int upperBound(int[] arr, int key) {  
        int lo = 0;  
        int hi = arr.length;

        while (lo < hi) {  
            int mid = lo + (hi - lo) / 2;  

            if (key < arr[mid]) {  
                hi = mid;  
            } else {  
                lo = mid+1;  
            }  
        }  
        return lo;  
    }  
}

 

---

 

lowerBound와 upperBound의 차이점은, arr[mid]==key일 때 lo와 hi 중 어느 것이 움직이는 가에 따라 달라진다. 즉, 이렇게 한 번에 구현하는 것도 가능하다.

flag가 true면 upperBound, false라면 lowerBound이다.

/**
 * @param flag : upperBound Flag
 */
private static int binarySearch(int[] arr, int key, boolean flag) {
    int lo = 0;
    int hi = arr.length;

    while (lo < hi) {
        int mid = lo + (hi - lo) / 2;

        if(key == arr[mid]) {
            if (flag) {
                lo = mid+1;
            } else {
                hi = mid;
            }
        }
        else if (key < arr[mid]) {
            hi = mid;
        } else if(key > arr[mid]) {
            lo = mid+1;
        }
    }
    return lo;
}