Question

855. 考场就座

English Version

题目描述

在考场里，一排有 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 上。

解法

方法一：有序集合 + 哈希表

考虑到每次 $seat()$ 时都需要找到最大距离的座位，我们可以使用有序集合来保存座位区间。有序集合的每个元素为一个二元组 $(l, r)$，表示 $l$ 和 $r$ 之间（不包括 $l$ 和 $r$）的座位可以坐学生。初始时有序集合中只有一个元素 $(-1, n)$，表示 $(-1, n)$ 之间的座位可以坐学生。

另外，我们使用两个哈希表 left 和 right 来维护每个有学生的座位的左右邻居学生，方便我们在 $leave(p)$ 时合并两个座位区间。

时间复杂度 $O(\log n)$，空间复杂度 $O(n)$。其中 $n$ 为考场的座位数。

from sortedcontainers import SortedList


class ExamRoom:
  def __init__(self, n: int):
    def dist(x):
      l, r = x
      return r - l - 1 if l == -1 or r == n else (r - l) >> 1

    self.n = n
    self.ts = SortedList(key=lambda x: (-dist(x), x[0]))
    self.left = {}
    self.right = {}
    self.add((-1, n))

  def seat(self) -> int:
    s = self.ts[0]
    p = (s[0] + s[1]) >> 1
    if s[0] == -1:
      p = 0
    elif s[1] == self.n:
      p = self.n - 1
    self.delete(s)
    self.add((s[0], p))
    self.add((p, s[1]))
    return p

  def leave(self, p: int) -> None:
    l, r = self.left[p], self.right[p]
    self.delete((l, p))
    self.delete((p, r))
    self.add((l, r))

  def add(self, s):
    self.ts.add(s)
    self.left[s[1]] = s[0]
    self.right[s[0]] = s[1]

  def delete(self, s):
    self.ts.remove(s)
    self.left.pop(s[1])
    self.right.pop(s[0])


# Your ExamRoom object will be instantiated and called as such:
# obj = ExamRoom(n)
# param_1 = obj.seat()
# obj.leave(p)

class ExamRoom {
  private TreeSet<int[]> ts = new TreeSet<>((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<>();
  private int n;

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

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

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

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

  private void add(int[] s) {
    ts.add(s);
    left.put(s[1], s[0]);
    right.put(s[0], s[1]);
  }

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

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

int N;

int dist(const pair<int, int>& p) {
  auto [l, r] = p;
  if (l == -1 || r == N) return r - l - 1;
  return (r - l) >> 1;
}

struct cmp {
  bool operator()(const pair<int, int>& a, const pair<int, int>& b) const {
    int d1 = dist(a), d2 = dist(b);
    return d1 == d2 ? a.first < b.first : d1 > d2;
  };
};

class ExamRoom {
public:
  ExamRoom(int n) {
    N = n;
    this->n = n;
    add({-1, n});
  }

  int seat() {
    auto s = *ts.begin();
    int p = (s.first + s.second) >> 1;
    if (s.first == -1) {
      p = 0;
    } else if (s.second == n) {
      p = n - 1;
    }
    del(s);
    add({s.first, p});
    add({p, s.second});
    return p;
  }

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

private:
  set<pair<int, int>, cmp> ts;
  unordered_map<int, int> left;
  unordered_map<int, int> right;
  int n;

  void add(pair<int, int> s) {
    ts.insert(s);
    left[s.second] = s.first;
    right[s.first] = s.second;
  }

  void del(pair<int, int> s) {
    ts.erase(s);
    left.erase(s.second);
    right.erase(s.first);
  }
};

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

type ExamRoom struct {
  rbt   *redblacktree.Tree
  left  map[int]int
  right map[int]int
  n   int
}

func Constructor(n int) ExamRoom {
  dist := func(s []int) int {
    if s[0] == -1 || s[1] == n {
      return s[1] - s[0] - 1
    }
    return (s[1] - s[0]) >> 1
  }
  cmp := func(a, b any) int {
    x, y := a.([]int), b.([]int)
    d1, d2 := dist(x), dist(y)
    if d1 == d2 {
      return x[0] - y[0]
    }
    return d2 - d1
  }
  this := ExamRoom{redblacktree.NewWith(cmp), map[int]int{}, map[int]int{}, n}
  this.add([]int{-1, n})
  return this
}

func (this *ExamRoom) Seat() int {
  s := this.rbt.Left().Key.([]int)
  p := (s[0] + s[1]) >> 1
  if s[0] == -1 {
    p = 0
  } else if s[1] == this.n {
    p = this.n - 1
  }
  this.del(s)
  this.add([]int{s[0], p})
  this.add([]int{p, s[1]})
  return p
}

func (this *ExamRoom) Leave(p int) {
  l, _ := this.left[p]
  r, _ := this.right[p]
  this.del([]int{l, p})
  this.del([]int{p, r})
  this.add([]int{l, r})
}

func (this *ExamRoom) add(s []int) {
  this.rbt.Put(s, struct{}{})
  this.left[s[1]] = s[0]
  this.right[s[0]] = s[1]
}

func (this *ExamRoom) del(s []int) {
  this.rbt.Remove(s)
  delete(this.left, s[1])
  delete(this.right, s[0])
}

/**
 * Your ExamRoom object will be instantiated and called as such:
 * obj := Constructor(n);
 * param_1 := obj.Seat();
 * obj.Leave(p);
 */