398. Random Pick Index Medium

1/**
2 * [398] Random Pick Index
3 *
4 * Given an integer array nums with possible duplicates, randomly output the index of a given target number. You can assume that the given target number must exist in the array.
5 * Implement the Solution class:
6 *
7 * Solution(int[] nums) Initializes the object with the array nums.
8 * int pick(int target) Picks a random index i from nums where nums[i] == target. If there are multiple valid i's, then each index should have an equal probability of returning.
9 *
10 *  
11 * Example 1:
12 *
13 * Input
14 * ["Solution", "pick", "pick", "pick"]
15 * [[[1, 2, 3, 3, 3]], [3], [1], [3]]
16 * Output
17 * [null, 4, 0, 2]
18 * Explanation
19 * Solution solution = new Solution([1, 2, 3, 3, 3]);
20 * solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
21 * solution.pick(1); // It should return 0. Since in the array only nums[0] is equal to 1.
22 * solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
23 *
24 *  
25 * Constraints:
26 *
27 * 1 <= nums.length <= 2 * 10^4
28 * -2^31 <= nums[i] <= 2^31 - 1
29 * target is an integer from nums.
30 * At most 10^4 calls will be made to pick.
31 *
32 */
33use std::collections::HashMap;
34pub struct Solution {
35    map: HashMap<i32, Vec<i32>>,
36}
37
38// problem: https://leetcode.com/problems/random-pick-index/
39// discuss: https://leetcode.com/problems/random-pick-index/discuss/?currentPage=1&orderBy=most_votes&query=
40
41// submission codes start here
42
43/**
44 * `&self` means the method takes an immutable reference.
45 * If you need a mutable reference, change it to `&mut self` instead.
46 */
47use rand::seq::SliceRandom;
48impl Solution {
49    fn new(nums: Vec<i32>) -> Self {
50        let m = nums.into_iter().enumerate().fold(
51            HashMap::new(),
52            |mut acc: HashMap<i32, Vec<i32>>, n| {
53                acc.entry(n.1)
54                    .and_modify(|x| {
55                        x.push(n.0 as i32);
56                    })
57                    .or_insert(vec![n.0 as i32]);
58                acc
59            },
60        );
61        Solution { map: m }
62    }
63
64    fn pick(&self, target: i32) -> i32 {
65        match self.map.get(&target) {
66            None => -1,
67            Some(x) => *x.choose(&mut rand::thread_rng()).unwrap(),
68        }
69    }
70}
71
72/**
73 * Your Solution object will be instantiated and called as such:
74 * let obj = Solution::new(nums);
75 * let ret_1: i32 = obj.pick(target);
76 */
77
78// submission codes end
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_398() {
86        let obj = Solution::new(vec![1,2,3,4,4,3,1]);
87        println!("Got {}", obj.pick(4));
88    }
89}
90