백준 15970 자바 - 화살표 그리기 (BOJ 15970 JAVA)

문제 : boj15970

 

  제시된 x와 y를 색상을 기준으로 나눠서 생각해보자. 즉, 문제에서 제시된 예시는 다음과 같이 나눠서 볼 수 있다. 두 경우는 연관되는 경우가 없으므로 색상이 다르다면 따로 생각해도 상관없다.

 

  이렇게 보면, 각 색상별로 각 지점에서 좌, 우로 가장 가까운 점을 선택하기만 하면 된다. 정렬을 한 후 이분탐색으로 쉽게 찾을 수 있다. 그럼 이럴 때 쓰기 딱 좋은 자료구조가 있는데 Binary Search Tree로 구현되어 있는 TreeSet이다. 일반적으로 set은 해당 값이 들어있는지 파악하는 용도로 많이 쓰지만(hashSet 처럼), TreeSet은 이분탐색트리로 구성되어 있고 알아서 정렬되어 들어가므로 ceilling과 floor도 쉽게 구할 수 있다. 즉, 그보다 높거나 낮은 element를 O(logN)에 알 수 있다.

 

  그럼 색상별로 TreeSet 자료구조를 두고 모든 점을 순회하며 해당 지점보다 높거나 낮은 위치의 element 중 가까운 지점을 선택해서 거리를 더해나가면 쉽게 답을 구할 수 있다. 예를들어 c 지점보다 낮은 지점은 a, 높은 지점은 e 이고 e가 더 가까우므로 거리는 2 임을 알 수 있다.

 

  참고로 e처럼 더 큰 값이 없는 경우 null로 나오는데, 예외처리를 하기 귀찮으면 이하 코드처럼 문제에서 제시된 값보다 훨씬 큰 값을 TreeSet에 넣어버리면 된다. (내 경우엔 -100만과 +100만을 넣었다.)

 

 

코드 : github

import java.io.DataInputStream;
import java.io.IOException;
import java.util.TreeSet;

public class Main extends FastInput {
    class Pos {
        int x, y;
        public Pos(int x, int y) {
            this.x = x;
            this.y = y;
        }
    }

    private void solution() throws Exception {
        int n = nextInt();
        TreeSet<Integer>[] tsArr = new TreeSet[n+1];

        Pos[] posArr = new Pos[n];
        for (int i = 0; i < n; i++) {
            int x = nextInt();
            int y = nextInt();
            posArr[i] = new Pos(x, y);
            if (tsArr[y] == null) {
                tsArr[y] = new TreeSet<>();
                tsArr[y].add(-1000000);
                tsArr[y].add(1000000);
            }
            tsArr[y].add(x);
        }

        int sum = 0;
        for (Pos pos : posArr) {
            sum += Math.min( Math.abs(pos.x - tsArr[pos.y].higher(pos.x)), Math.abs(pos.x - tsArr[pos.y].lower(pos.x)));
        }
        System.out.println(sum);
    }

    public static void main(String[] args) throws Exception {
        initFI();
        new Main().solution();
    }
}

class FastInput {
    private static final int DEFAULT_BUFFER_SIZE = 1 << 16;
    private static DataInputStream inputStream;
    private static byte[] buffer;
    private static int curIdx, maxIdx;

    protected static void initFI() {
        inputStream = new DataInputStream(System.in);
        buffer = new byte[DEFAULT_BUFFER_SIZE];
        curIdx = maxIdx = 0;
    }

    protected static int nextInt() throws IOException {
        int ret = 0;
        byte c = read();
        while (c <= ' ') c = read();
        do {
            ret = ret * 10 + c - '0';
        } while ((c = read()) >= '0' && c <= '9');
        return ret;
    }

    private static byte read() throws IOException {
        if (curIdx == maxIdx) {
            maxIdx = inputStream.read(buffer, curIdx = 0, DEFAULT_BUFFER_SIZE);
            if (maxIdx == -1) buffer[0] = -1;
        }
        return buffer[curIdx++];
    }
}