leetcode 855. 考场就座

@(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);
 */