rustgym/leetcode/
_381_insert_delete_get_random_o1_duplicate_allowed.rs

1use rand::prelude::*;
2use std::collections::BinaryHeap;
3use std::collections::HashMap;
4
5#[derive(Debug)]
6struct RandomizedCollection {
7    rng: ThreadRng,
8    indexes: HashMap<i32, BinaryHeap<usize>>,
9    choices: Vec<i32>,
10}
11
12impl RandomizedCollection {
13    fn new() -> Self {
14        let rng = thread_rng();
15        let indexes = HashMap::new();
16        let choices = vec![];
17        RandomizedCollection {
18            rng,
19            indexes,
20            choices,
21        }
22    }
23
24    fn insert(&mut self, val: i32) -> bool {
25        let ids = self.indexes.entry(val).or_default();
26        ids.push(self.choices.len());
27        self.choices.push(val);
28        ids.len() == 1
29    }
30
31    fn remove(&mut self, val: i32) -> bool {
32        if let Some(ids) = self.indexes.get_mut(&val) {
33            if let Some(index) = ids.pop() {
34                let last_index = self.choices.len() - 1;
35                let last_value = self.choices[last_index];
36                if last_value != val {
37                    let other_ids = self.indexes.get_mut(&last_value).unwrap();
38                    other_ids.pop();
39                    other_ids.push(index);
40                    self.choices.swap(index, last_index);
41                }
42                self.choices.pop();
43                true
44            } else {
45                false
46            }
47        } else {
48            false
49        }
50    }
51
52    fn get_random(&mut self) -> i32 {
53        self.choices[self.rng.gen_range(0, self.choices.len())]
54    }
55}
56
57#[test]
58fn test() {
59    let mut obj = RandomizedCollection::new();
60    assert_eq!(obj.remove(0), false);
61    dbg!(&obj);
62    // assert_eq!(obj.insert(1), false);
63    // dbg!(&obj);
64    // assert_eq!(obj.insert(2), true);
65    // dbg!(&obj);
66    // assert_eq!(obj.insert(2), false);
67    // dbg!(&obj);
68    // assert_eq!(obj.insert(2), false);
69    // dbg!(&obj);
70    // assert_eq!(obj.remove(1), true);
71    // dbg!(&obj);
72    // assert_eq!(obj.remove(1), true);
73    // dbg!(&obj);
74    // assert_eq!(obj.remove(2), true);
75    // dbg!(&obj);
76    // assert_eq!(obj.insert(1), true);
77    // dbg!(&obj);
78    // assert_eq!(obj.remove(2), true);
79}
rustgym/leetcode/_381_insert_delete_get_random_o1_duplicate_allowed.rs

rustgym/leetcode/
_381_insert_delete_get_random_o1_duplicate_allowed.rs
