diff --git a/rust/0380-insert-delete-getrandom-o1.rs b/rust/0380-insert-delete-getrandom-o1.rs new file mode 100644 index 000000000..61913abb1 --- /dev/null +++ b/rust/0380-insert-delete-getrandom-o1.rs @@ -0,0 +1,40 @@ +use rand::seq::SliceRandom; +use std::collections::HashMap; + +pub struct RandomizedSet { + mp: HashMap, + arr: Vec, +} + +impl RandomizedSet { + fn new() -> Self { + RandomizedSet { + mp: HashMap::new(), + arr: Vec::new(), + } + } + + fn insert(&mut self, val: i32) -> bool { + let res = !self.mp.contains_key(&val); + if res { + self.mp.insert(val, self.arr.len() as i32); + self.arr.push(val); + } + res + } + fn remove(&mut self, val: i32) -> bool { + let res = self.mp.contains_key(&val); + if res { + let idx = *self.mp.get(&val).unwrap(); + self.mp + .entry(*self.arr.last().unwrap()) + .and_modify(|v| *v = idx); + self.arr.swap_remove(idx as usize); + self.mp.remove(&val); + } + res + } + fn get_random(&self) -> i32 { + *self.arr.choose(&mut rand::thread_rng()).unwrap() + } +}