leetcode398. Random Pick Index
題目要求
Given an array of integers 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. Note: The array size can be very large. Solution that uses too much extra space will not pass the judge. Example: int[] nums = new int[] {1,2,3,3,3}; Solution solution = new Solution(nums); // pick(3) should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning. solution.pick(3); // pick(1) should return 0. Since in the array only nums[0] is equal to 1. solution.pick(1);
設計一個數據結構,使得從該資料結構中查詢一個數字時,能夠以等概率返回該數字所在的任何下標。額外的要求是隻要佔用O(1)的額外的空間複雜度。
思路一:O(2n)的實現
其實要是想在O(1)的時間內完成隨機數的獲取,只需要快取每個數字出現的下標,但是這意味著需要先對資料進行遍歷,並且還需要O(n)的空間來額外儲存數字下標的集合。這違背了題目意思。所以我們只能每次在獲取數字的時候來統計數字出現的次數,然後針對次數獲取隨機數下標。
程式碼如下:
private int[] nums; private Random r; public Solution(int[] nums) { this.nums = nums; this.r = new Random(); } public int pick(int target) { int count = 0; int result = -1; for(int i = 0 ; i<nums.length ; i++) { if(target == nums[i]) { count++; } } int index = r.nextInt(count); for(int i = 0 ; i<nums.length ; i++) { if(target == nums[i]) { index--; if(index == -1){ result = i; break; } } } return result; }
思路二:蓄水池問題
有沒有可能在一次的遍歷中,找到這個隨機下標呢?這就涉及到一個概率問題,即當我在遍歷過程中遇到這個數字時,我是否需要選擇它作為我的結果值。
首先介紹一下蓄水池抽樣演算法 。蓄水池抽樣演算法主要對應的是這樣的一個問題,假設有一個非常大的資料集,或者是一個以流的形式作為輸入的資料集,希望從中選擇K個樣本,此時我們可能無法把所有的樣本都放在記憶體中,因此只能儲存一個大小為K的集合,每次遇到一個數據時,決定是否將其放入樣本中。
因此,假設當前遇到的資料總共有N個,如果N小於K,則每個資料被選中的概率為1,如果N>K,則每個資料被選中的概率為K/N,舊樣本中資料被保留的概率為1 - K(N-K)/N
即K/N,因此每一箇舊的資料被替換掉的概率為1/N。
按照這個思路,我們在下面的遍歷時,每遇到一個新元素,就判斷當前選中的元素是否會被該元素替換掉即可。
private int[] nums; private Random r; public Solution(int[] nums) { this.nums = nums; this.r = new Random(); } public int pick(int target) { int count = 0; int result = -1; for(int i = 0 ; i<nums.length ; i++) { if(nums[i] == target) { int index = r.nextInt(++count); if(index == 0) { result = i; } } } return result; }