LeetCode/src/main/java/com/wzx/leetcode/No855ExamRoom.java at master · wzx140/LeetCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package com.wzx.leetcode;

import java.util.*;

/**
 * @see <a href="https://leetcode.com/problems/exam-room/">https://leetcode.com/problems/exam-room/</a>
 * @author wzx
 */
public class No855ExamRoom {

  /**
   * 将所有考生看成两两成对的线段
   */
  public static class ExamRoom {

    /**
     * 根据自定义长度排序的线段集合
     */
    private final TreeSet<int[]> intervalSet;

    /**
     * 以key为左端点的线段
     */
    private final Map<Integer, int[]> leftMap;

    /**
     * 以key为右端点的线段
     */
    private final Map<Integer, int[]> rightMap;

    private final int N;

    /**
     * 返回线段端点(exclude)至中点(include)的距离
     */
    private int distance(int[] interval) {
      int left = interval[0];
      int right = interval[1];
      if (left == -1) return right;
      if (right == N) return N - left - 1;
      return (right - left) / 2;
    }

    private void remove(int[] interval){
      intervalSet.remove(interval);
      leftMap.remove(interval[0]);
      rightMap.remove(interval[1]);
    }

    private void add(int[] interval){
      intervalSet.add(interval);
      leftMap.put(interval[0], interval);
      rightMap.put(interval[1], interval);
    }

    public ExamRoom(int N) {
      this.N = N;
      intervalSet = new TreeSet<>(
              (interval1, interval2) -> {
                int dist1 = distance(interval1);
                int dist2 = distance(interval2);
                // 优先分配索引较小的位置
                if (dist1 == dist2) {
                  return interval1[0] - interval2[0];
                }
                return dist2 - dist1;
              });
      leftMap = new HashMap<>();
      rightMap = new HashMap<>();
      // dummy
      int[] dummyInterval = new int[]{-1, N};
      intervalSet.add(dummyInterval);
      leftMap.put(-1, dummyInterval);
      rightMap.put(N, dummyInterval);
    }

    public int seat() {
      // 取出最大距离线段
      int[] interval = intervalSet.first();
      // 分裂成两个
      int seat;
      if (interval[0] == -1) {
        seat = 0;
      } else if (interval[1] == N) {
        seat = N - 1;
      } else {
        seat = (interval[1] - interval[0]) / 2 + interval[0];
      }
      int[] leftInterval = new int[]{interval[0], seat};
      int[] rightInterval = new int[]{seat, interval[1]};
      remove(interval);
      add(leftInterval);
      add(rightInterval);

      return seat;
    }

    public void leave(int p) {
      // 合并此位置两侧的线段
      int[] leftInterval = rightMap.get(p);
      int[] rightInterval = leftMap.get(p);
      int[] newInterval = new int[]{leftInterval[0], rightInterval[1]};
      remove(leftInterval);
      remove(rightInterval);
      add(newInterval);
    }
  }
}