rustgym/leetcode/
_855_exam_room.rs

1use std::cmp::Reverse;
2use std::collections::BTreeSet;
3use std::collections::HashMap;
4
5// (distance, left, right)
6type Segment = (Reverse<i32>, i32, i32);
7
8#[derive(Debug)]
9struct ExamRoom {
10    n: i32,
11    segments: BTreeSet<Segment>,
12    l_indexes: HashMap<i32, i32>,
13    r_indexes: HashMap<i32, i32>,
14}
15
16impl ExamRoom {
17    fn new(n: i32) -> Self {
18        let mut segments = BTreeSet::new();
19        segments.insert(Self::segment(0, n - 1, n));
20        let mut l_indexes = HashMap::new();
21        let mut r_indexes = HashMap::new();
22        l_indexes.insert(0, n - 1);
23        r_indexes.insert(n - 1, 0);
24        ExamRoom {
25            n,
26            segments,
27            l_indexes,
28            r_indexes,
29        }
30    }
31
32    fn seat(&mut self) -> i32 {
33        let mut it = self.segments.iter();
34        if let Some(&first) = it.next() {
35            let l = first.1;
36            let r = first.2;
37            let p = Self::split(&first, self.n);
38            self.segments.remove(&first);
39            self.segments.insert(Self::segment(l, p - 1, self.n));
40            self.segments.insert(Self::segment(p + 1, r, self.n));
41            self.l_indexes.insert(l, p - 1);
42            self.r_indexes.insert(p - 1, l);
43            self.l_indexes.insert(p + 1, r);
44            self.r_indexes.insert(r, p + 1);
45            p
46        } else {
47            -1
48        }
49    }
50
51    fn leave(&mut self, p: i32) {
52        let r1 = p - 1;
53        let l1 = self.r_indexes[&r1];
54        let l2 = p + 1;
55        let r2 = self.l_indexes[&l2];
56        self.segments.remove(&Self::segment(l1, r1, self.n));
57        self.segments.remove(&Self::segment(l2, r2, self.n));
58        self.segments.insert(Self::segment(l1, r2, self.n));
59        self.r_indexes.remove(&r1);
60        self.l_indexes.remove(&l2);
61        self.l_indexes.insert(l1, r2);
62        self.r_indexes.insert(r2, l1);
63    }
64
65    fn segment(l: i32, r: i32, n: i32) -> Segment {
66        if l == 0 {
67            return (Reverse(r), l, r);
68        }
69        if r == n - 1 {
70            return (Reverse(n - 1 - l), l, r);
71        }
72        if l <= r {
73            (Reverse((r - l) / 2), l, r)
74        } else {
75            (Reverse(-1), l, r)
76        }
77    }
78
79    fn split(s: &Segment, n: i32) -> i32 {
80        let l = s.1;
81        let r = s.2;
82        if l == 0 {
83            return 0;
84        }
85        if r == n - 1 {
86            return n - 1;
87        }
88        l + (r - l) / 2
89    }
90}
91
92#[test]
93fn test() {
94    let mut exam_room = ExamRoom::new(10);
95    assert_eq!(exam_room.seat(), 0);
96    assert_eq!(exam_room.seat(), 9);
97    assert_eq!(exam_room.seat(), 4);
98    assert_eq!(exam_room.seat(), 2);
99    exam_room.leave(4);
100    assert_eq!(exam_room.seat(), 5);
101
102    let mut exam_room = ExamRoom::new(4);
103    assert_eq!(exam_room.seat(), 0);
104    assert_eq!(exam_room.seat(), 3);
105    assert_eq!(exam_room.seat(), 1);
106    assert_eq!(exam_room.seat(), 2);
107    exam_room.leave(1);
108    exam_room.leave(3);
109    assert_eq!(exam_room.seat(), 1);
110}
rustgym/leetcode/_855_exam_room.rs

rustgym/leetcode/
_855_exam_room.rs
