@(labuladong的算法小抄)[平衡二叉搜尋樹]

leetcode 855. 考場就座

題目描述

解題思路

參考：labuladong的算法小抄P389

TreeSet是一種有序結構，底層是由紅黑樹（一種平衡二叉搜尋樹）維護有序性。并且查找、增加、删除節點的操作複雜度都是o(logn).

class ExamRoom {
    /* p -> 以p為左端點的線段 */
    private Map<Integer, int[]> startMap;
    /* p -> 以p為右端點的線段 */
    private Map<Integer, int[]> endMap;
    /* 根據[x,y]的線段長度從小到大存放所有線段 */
    private TreeSet<int[]> xy;
    private int N;

    public ExamRoom(int N) {
        this.N = N;
        startMap = new HashMap<>();
        endMap = new HashMap<>();
        xy = new TreeSet<>((a, b) -> {
            /* 算出兩個線段的中點到端點的長度 */
            int distA = distance(a);
            int distB = distance(b);
            /* 如果長度相同，則按索引降序，確定排在後面的索引小一些 */
            if (distA == distB) {
                return b[0] - a[0];
            }
            return distA - distB;
        });
        /* 在有序集合中先放一個虛拟線段 */
        addInterval(new int[]{-1, N});
    }

    /* 封裝所有 增加線段 的操作 */
    private void addInterval(int[] intv) {
        xy.add(intv);
        startMap.put(intv[0], intv);
        endMap.put(intv[1], intv);
    }

    /* 封裝所有 删除線段 的操作 */
    private void removeInterval(int[] intv) {
        xy.remove(intv);
        startMap.remove(intv[0]);
        endMap.remove(intv[1]);
    }

    /* 計算線段的中點到端點的長度 */
    private int distance(int[] intv) {
        int x = intv[0];
        int y = intv[1];
        /* 中點是0 */
        if (x == -1) return y;
        /* 中點是N-1 */
        if (y == N) return N - 1 - x;
        /* 中點和端點之間的長度 */
        return (y - x) / 2;
    }
    /* 給學生配置設定一個位置，傳回該位置 */
    public int seat() {
        /* 從有序集合中取出最長的線段 */
        int[] longest = xy.last();
        int x = longest[0];
        int y = longest[1];
        int seat;
        if (x == -1) {
            /* 情況一：左邊沒人 */
            seat = 0;
        } else if (y == N) {
            /* 情況二：右邊沒人 */
            seat = N - 1;
        } else {
            /* 情況三，不是邊界的話，就坐中間 */
            seat = x + (y - x) / 2;
        }
        /* 将最長的線段分成左右兩段 */
        int[] left = new int[]{x, seat};
        int[] right = new int[]{seat, y};
        removeInterval(longest);
        addInterval(left);
        addInterval(right);
        return seat;
    }

    public void leave(int p) {
        /* 找到p對應的右邊和左邊線段 */
        int[] right = startMap.get(p);
        int[] left = endMap.get(p);
        /* 将兩條線段合并成一條線段 */
        int[] mergedIntv = new int[]{left[0], right[1]};
        /* 删除舊線段，插入新線段 */
        removeInterval(left);
        removeInterval(right);
        addInterval(mergedIntv);
    }
}

/**
 * Your ExamRoom object will be instantiated and called as such:
 * ExamRoom obj = new ExamRoom(N);
 * int param_1 = obj.seat();
 * obj.leave(p);
 */