Exam Room LeetCode solution in Java

1
class ExamRoom {
2

3
  private final int max;
4
  private final TreeSet<Interval> available;
5
  private final TreeSet<Integer> taken;
6

7
  public ExamRoom(int n) {
8
    this.max = n - 1;
9
    this.available =
10
        new TreeSet<>(
11
            (a, b) -> {
12
              var distA = getMinDistance(a);
13
              var distB = getMinDistance(b);
14
              return distA == distB ? a.s - b.s : distB - distA;
15
            });
16
    this.available.add(new Interval(0, max));
17
    this.taken = new TreeSet<>();
18
  }
19

20
  public int seat() {
21
    var inter = available.pollFirst();
22
    var idx = getInsertPosition(inter);
23
    taken.add(idx);
24
    if ((idx - 1) - inter.s >= 0) available.add(new Interval(inter.s, idx - 1));
25
    if (inter.e - (idx + 1) >= 0) available.add(new Interval(idx + 1, inter.e));
26
    return idx;
27
  }
28

29
  public void leave(int p) {
30
    taken.remove(p);
31
    var lo = taken.lower(p);
32
    if (lo == null) lo = -1;
33
    var hi = taken.higher(p);
34
    if (hi == null) hi = max + 1;
35
    available.remove(new Interval(lo + 1, p - 1));
36
    available.remove(new Interval(p + 1, hi - 1));
37
    available.add(new Interval(lo + 1, hi - 1));
38
  }
39

40
  private int getInsertPosition(Interval inter) {
41
    if (inter.s == 0) return 0;
42
    else if (inter.e == max) return max;
43
    else return inter.s + (inter.e - inter.s) / 2;
44
  }
45

46
  private int getMinDistance(Interval in) {
47
    return in.s == 0 || in.e == max ? in.e - in.s : (in.e - in.s) / 2;
48
  }
49

50
  private final class Interval {
51
    private final int s;
52
    private final int e;
53

54
    Interval(int s, int e) {
55
      this.s = s;
56
      this.e = e;
57
    }
58

59
    @Override
60
    public String toString() {
61
      return "[" + s + "," + e + "]";
62
    }
63
  }
64
}