考场就座

标签: 设计 有序集合 堆(优先队列)

难度: Medium

在考场里,一排有 N 个座位,分别编号为 0, 1, 2, ..., N-1 。

当学生进入考场后,他必须坐在能够使他与离他最近的人之间的距离达到最大化的座位上。如果有多个这样的座位,他会坐在编号最小的座位上。(另外,如果考场里没有人,那么学生就坐在 0 号座位上。)

返回 ExamRoom(int N) 类,它有两个公开的函数:其中,函数 ExamRoom.seat() 会返回一个 int (整型数据),代表学生坐的位置;函数 ExamRoom.leave(int p) 代表坐在座位 p 上的学生现在离开了考场。每次调用 ExamRoom.leave(p) 时都保证有学生坐在座位 p 上。

示例:

输入:["ExamRoom","seat","seat","seat","seat","leave","seat"], [[10],[],[],[],[],[4],[]]
输出:[null,0,9,4,2,null,5]
解释:
ExamRoom(10) -> null
seat() -> 0,没有人在考场里,那么学生坐在 0 号座位上。
seat() -> 9,学生最后坐在 9 号座位上。
seat() -> 4,学生最后坐在 4 号座位上。
seat() -> 2,学生最后坐在 2 号座位上。
leave(4) -> null
seat() -> 5,学生最后坐在 5 号座位上。

提示:

  1. 1 <= N <= 10^9
  2. 在所有的测试样例中 ExamRoom.seat() 和 ExamRoom.leave() 最多被调用 10^4 次。
  3. 保证在调用 ExamRoom.leave(p) 时有学生正坐在座位 p 上。

Submission

运行时间: 64 ms

内存: 0.0 MB

from heapq import heappop, heappush


class ExamRoom(object):

    def __init__(self, N):
        """
        :type N: int
        """
        self.N = N
        self.heap = []
        self.avail_first = {}
        self.avail_last = {}
        self.put_segment(0, self.N - 1)

    def put_segment(self, first, last):

        if first == 0 or last == self.N - 1:
            priority = last - first
        else:
            priority = (last - first) // 2

        segment = [-priority, first, last, True]

        self.avail_first[first] = segment
        self.avail_last[last] = segment

        heappush(self.heap, segment)

    def seat(self):
        """
        :rtype: int
        """
        while True:
            _, first, last, is_valid = heappop(self.heap)

            if is_valid:
                del self.avail_first[first]
                del self.avail_last[last]
                break

        if first == 0:
            ret = 0
            if first != last:
                self.put_segment(first + 1, last)

        elif last == self.N - 1:
            ret = last
            if first != last:
                self.put_segment(first, last - 1)

        else:
            ret = first + (last - first) // 2

            if ret > first:
                self.put_segment(first, ret - 1)

            if ret < last:
                self.put_segment(ret + 1, last)

        return ret

    def leave(self, p):
        """
        :type p: int
        :rtype: void
        """
        first = p
        last = p

        left = p - 1
        right = p + 1

        if left >= 0 and left in self.avail_last:
            segment_left = self.avail_last.pop(left)
            segment_left[3] = False
            first = segment_left[1]

        if right < self.N and right in self.avail_first:
            segment_right = self.avail_first.pop(right)
            segment_right[3] = False
            last = segment_right[2]

        self.put_segment(first, last)

Explain

这个题解使用了优先队列(最大堆)来快速找到最优的座位。首先,每个区间(即连续的空座位)被标记为一个三元组 [-priority, first, last],其中 priority 是该区间的优先级,first 和 last 分别是区间的起始和结束座位。如果区间是开头或结尾的座位,优先级就是区间长度;否则,优先级是区间长度除以二的商。这样设计是为了优先坐在最中间的位置,以最大化与最近的人的距离。当座位被选取时,相关区间会被拆分并重新加入堆中,如果座位 p 被释放,相邻的空座位区间会合并并更新。

时间复杂度: O(log n)

空间复杂度: O(n)

from heapq import heappop, heappush


class ExamRoom(object):

    def __init__(self, N):
        self.N = N  # 总座位数
        self.heap = []  # 用于维护所有区间的最大堆
        self.avail_first = {}  # 存储以区间起点为键的区间映射
        self.avail_last = {}  # 存储以区间终点为键的区间映射
        self.put_segment(0, self.N - 1)  # 初始化,一开始整个区间都是空的

    def put_segment(self, first, last):
        if first == 0 or last == self.N - 1:
            priority = last - first  # 如果区间是边界,则取整个区间长度作为优先级
        else:
            priority = (last - first) // 2  # 否则取中间位置为优先级
        segment = [-priority, first, last, True]  # 创建区间,负号确保优先队列为最大堆
        self.avail_first[first] = segment  # 记录区间起点
        self.avail_last[last] = segment  # 记录区间终点
        heappush(self.heap, segment)  # 将区间推入堆中

    def seat(self):
        while True:
            _, first, last, is_valid = heappop(self.heap)  # 取出优先级最高的区间
            if is_valid:
                del self.avail_first[first]  # 删除旧的映射
                del self.avail_last[last]  # 删除旧的映射
                break  # 找到有效区间,跳出循环
        if first == 0:
            ret = 0  # 选择第一个座位
            if first != last:
                self.put_segment(first + 1, last)  # 更新区间
        elif last == self.N - 1:
            ret = last  # 选择最后一个座位
            if first != last:
                self.put_segment(first, last - 1)  # 更新区间
        else:
            ret = first + (last - first) // 2  # 选择中间座位
            if ret > first:
                self.put_segment(first, ret - 1)  # 更新区间
            if ret < last:
                self.put_segment(ret + 1, last)  # 更新区间
        return ret  # 返回选定的座位

    def leave(self, p):
        first = p  # 开始点是 p
        last = p  # 结束点也是 p
        left = p - 1  # 检查左邻座位
        right = p + 1  # 检查右邻座位
        if left >= 0 and left in self.avail_last:
            segment_left = self.avail_last.pop(left)  # 获取并删除左边的区间
            segment_left[3] = False  # 标记为无效
            first = segment_left[1]  # 扩展合并区间的起点
        if right < self.N and right in self.avail_first:
            segment_right = self.avail_first.pop(right)  # 获取并删除右边的区间
            segment_right[3] = False  # 标记为无效
            last = segment_right[2]  # 扩展合并区间的终点
        self.put_segment(first, last)  # 合并区间并重新放入堆中

Explore

是的,初始化ExamRoom时整个座位区间设定为从0到N-1表示开始时考场是完全空的。这样的设定是为了模拟考场在开始前没有任何学生入座的情况,从而允许算法从一个完整且未被占用的座位区间开始进行座位分配。这种方式简化了座位的管理和分配过程,因为它从最大可能的空座位区间开始,并逐步根据实际座位的选择和释放进行调整。

在put_segment函数中,边界区间(即第一个或最后一个座位的区间)的优先级设定为整个区间长度,而非边界区间(即中间的座位区间)的优先级设置为区间长度除以二的原因在于优化座位的选择,使人与人之间的距离最大化。边界区间的优先级较高是因为在边界坐下时,一个方向上没有其他人,从而可以确保一侧的最大距离。对于中间的区间,通过取区间长度的一半作为优先级,可以确保选择座位时人们尽可能处于区间的中间位置,从而最大化与最近的两侧人的距离。这种处理有助于在考场中尽可能地避免紧挨着坐,提高考场的使用效率和考生的舒适度。

在seat函数中,确保从堆中弹出的区间是有效的通过在每个区间元组中添加一个布尔值来实现,该布尔值标记区间是否有效。当执行leave操作时,如果某个座位被释放,与该座位相邻的区间可能会被合并,这时原有的区间会被标记为无效并从有效映射中删除。从堆中弹出区间后,会检查这个标记来确定区间是否仍然有效。如果区间无效(即布尔值为False),则继续从堆中弹出下一个区间,直到找到一个有效的区间。这种方法确保了即使区间在堆中排序后状态发生变化,座位分配的逻辑仍然能够正确地处理有效区间,避免了无效数据的干扰。