LeetCode：855. 考场就座（TreeSet Java）_考场安排算法 java-优快云博客

855. 考场就座

题目描述：

在考场里，有 n 个座位排成一行，编号为 0 到 n - 1。

当学生进入考场后，他必须坐在离最近的人最远的座位上。如果有多个这样的座位，他会坐在编号最小的座位上。(另外，如果考场里没有人，那么学生就坐在 0 号座位上。)

设计一个模拟所述考场的类。

实现 ExamRoom 类：

ExamRoom(int n) 用座位的数量 n 初始化考场对象。
int seat() 返回下一个学生将会入座的座位编号。
void leave(int p) 指定坐在座位 p 的学生将离开教室。保证座位 p 上会有一位学生。

示例 1：

输入：
["ExamRoom", "seat", "seat", "seat", "seat", "leave", "seat"]
[[10], [], [], [], [], [4], []]
输出：
[null, 0, 9, 4, 2, null, 5]
解释：
ExamRoom examRoom = new ExamRoom(10);
examRoom.seat(); // 返回 0，房间里没有人，学生坐在 0 号座位。
examRoom.seat(); // 返回 9，学生最后坐在 9 号座位。
examRoom.seat(); // 返回 4，学生最后坐在 4 号座位。
examRoom.seat(); // 返回 2，学生最后坐在 2 号座位。
examRoom.leave(4);
examRoom.seat(); // 返回 5，学生最后坐在 5 号座位。

提示：

1 <= n <= 109
保证有学生正坐在座位 p 上。
seat 和 leave 最多被调用 104 次。

实现代码与解析：

TreeSet

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.TreeSet;

class ExamRoom {

    int n;
    private PriorityQueue<int[]> ts = new PriorityQueue<>((a, b) -> {
        int d1 = dist(a), d2 = dist(b);
        return d1 == d2 ? a[0] - b[0] : d2 - d1;
    });


    private Map<Integer, Integer> left = new HashMap<>();
    private Map<Integer, Integer> right = new HashMap<>();

    public ExamRoom(int n) {
        this.n = n;
        add(-1, n);
    }

    public int seat() {
        int[] peek = ts.peek();
        int p = (peek[0] + peek[1]) >> 1;
        if (peek[0] == -1) p = 0;
        else if (peek[1] == n) p = n - 1;
        del(peek);
        add(peek[0], p);
        add(p, peek[1]);
        return p;
    }

    public void leave(int p) {
        int l = left.get(p);
        int r = right.get(p);
        del(new int[]{l, p});
        del(new int[]{p, r});
        add(l, r);
    }

    public void add(int l, int r) {
        ts.add(new int[]{l, r});
        left.put(r, l);
        right.put(l, r);
    }

    public void del(int[] q) {
        ts.remove(q);
        left.remove(q[1]);
        right.remove(q[0]);
    }

    private int dist(int[] q) {
        int l = q[0], r = q[1];
        return l == -1 || r == n ? r - l - 1 : (r - l) >> 1;
    }

}

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

原理思路：

1. 重要的成员变量

int n：表示考场中座位的总数，限定了座位范围的边界值。

TreeSet<int[]> ts：这是一个自定义比较器的TreeSet，用来存储表示座位区间的数组。比较器的逻辑是先比较区间的 “距离”（通过dist方法计算，距离的定义与区间两端的情况有关），如果距离相等则按照区间左端点的大小来排序。其目的是方便快速找到最合适的空闲座位区间来分配座位。

Map<Integer, Integer> left和Map<Integer, Integer> right：这两个HashMap分别用于记录每个座位区间的左端点对应的右端点，以及右端点对应的左端点，方便在座位分配和释放操作时快速查找和更新区间的关联关系。

2. 构造函数

public ExamRoom(int n)：构造函数接收表示座位总数的参数n，初始化this.n为传入的n值，并且调用add(-1, n)方法，将初始的整个考场座位区间（从虚拟的-1到实际的n，代表初始全部空闲的状态）添加到数据结构中进行管理。

3. seat方法（座位分配方法）

首先获取ts中第一个（按照自定义比较器排序最合适的）座位区间数组peek。

然后计算要分配的座位位置p，计算规则是取区间的中间位置（通过(peek[0] + peek[1]) >> 1，也就是取平均值向下取整），同时针对区间左端点是-1（代表最左边开始的区间）或者右端点是n（代表最右边的区间）的特殊情况，进行边界处理，分别将座位分配到最左边（p = 0）或者最右边（p = n - 1）。

接着通过del方法移除当前选中的座位区间，再调用两次add方法，将因为分配座位而产生的两个新的区间添加到数据结构中，最后返回分配的座位位置p。

4. leave方法（座位释放方法）

接收要释放的座位位置参数p，通过left和right这两个Map获取该座位所在区间的左端点l和右端点r。

调用del方法移除包含该座位的两个区间（分别是[l, p]和[p, r]），然后调用add方法将合并后的区间[l, r]重新添加到数据结构中，表示座位释放后区间状态的更新。

5. add方法（添加区间方法）

接收区间的左端点l和右端点r作为参数，将代表区间的数组{l, r}添加到ts中，同时在left和right这两个Map中分别添加对应端点的关联关系，方便后续查找和更新操作。

6. del方法（移除区间方法）

接收表示区间的数组q作为参数，从ts中移除该区间，同时在left和right这两个Map中移除对应端点的关联记录，实现区间从数据结构中的彻底删除。

7. dist方法（计算区间距离方法）

接收表示区间的数组q，获取其左端点l和右端点r，根据区间两端是否是边界情况（l == -1或者r == n）来计算区间的 “距离”，如果是边界情况距离计算为r - l - 1，否则为(r - l) >> 1（即区间长度的一半），这个距离值用于在TreeSet的比较器中判断区间的优先级，从而确定分配座位时优先选择的区间。