update 398

Author: fishercoder1534

Diff for: src/main/java/com/fishercoder/solutions/_398.java

@@ -1,33 +1,28 @@
 package com.fishercoder.solutions;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 
 public class _398 {
 
     //TODO: use reservoir sampling to solve it again
 
     public static class Solution {
-        //brute force
-        int[] input;
-        java.util.Random rand = new java.util.Random();
+        Map<Integer, List<Integer>> map;
+        Random random;
 
         public Solution(int[] nums) {
-            input = nums;
+            map = new HashMap<>();
+            random = new Random();
+            for (int i = 0; i < nums.length; i++) {
+                List<Integer> list = map.getOrDefault(nums[i], new ArrayList<>());
+                list.add(i);
+                map.put(nums[i], list);
+            }
         }
 
         public int pick(int target) {
-            List<Integer> list = new ArrayList();
-            for (int i = 0; i < input.length; i++) {
-                if (input[i] == target) {
-                    list.add(i);
-                }
-            }
-            if (list.size() == 1) {
-                return list.get(0);
-            }
-            int randomIndex = rand.nextInt(list.size());
-            return list.get(randomIndex);
+            List<Integer> list = map.get(target);
+            return list.get(random.nextInt(list.size()));
         }
     }
 }