题目描述

[EN | CN]

给定一个可能含有重复元素的整数数组,要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。

注意:

数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。

示例:

1
2
3
4
5
6
7
8
int[] nums = new int[] {1,2,3,3,3};
Solution solution = new Solution(nums);

// pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
solution.pick(3);

// pick(1) 应该返回 0。因为只有nums[0]等于1。
solution.pick(1);

解法 1:收集再随机输出

比较直接的思路就是收集所有可选的索引,然后随机输出一个。

假设数组长度为 $n$,符合条件的元素数量为 $m$,那么

  • 时间复杂度为 $O(n)$,遍历一遍;
  • 空间复杂度为 $O(m)$,需要存储符合条件的所有索引,如果 $m$ 太大,其实也是不行的。

实现与结果如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
import random

class Solution:

    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        candidates = [idx for idx, i in enumerate(self.nums) if i == target]
        return random.choice(candidates)


# Your Solution object will be instantiated and called as such:
# obj = Solution(nums)
# param_1 = obj.pick(target)xxxxxxxx
  • 执行用时:476 ms,在所有 Python3 提交中击败了 21.94% 的用户。
  • 内存消耗:17.5 MB,在所有 Python3 提交中击败了 50.00% 的用户。

解法 2:蓄水池算法

关于蓄水池算法的介绍可以参考《LeetCode 382. 链表随机结点 ⌨️》,本题相当于那道题中对链表进行一次筛选(是否为目标值)再选。

该算法在本题中可以很顺畅地直接使用,而且可以将空间复杂度降为 $O(1)$。

假设数组长度为 $n$,符合条件的元素数量为 $m$,那么

  • 时间复杂度为 $O(n)$,遍历一遍;
  • 空间复杂度为 $O(1)$,不需要存储额外内容。

实现与结果如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import random

class Solution:

    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        ret = -1
        for i, num in enumerate(self.nums):
            if num == target:
                ret = i
                break
        count = 1

        for i in range(ret + 1, len(self.nums)):
            if self.nums[i] == target:
                count += 1
                r = random.random()
                if r < 1 / count:
                    ret = i

        return ret

# Your Solution object will be instantiated and called as such:
# obj = Solution(nums)
# param_1 = obj.pick(target)
  • 执行用时:352 ms,在所有 Python3 提交中击败了 76.45% 的用户。
  • 内存消耗:17.3 MB,在所有 Python3 提交中击败了 50.00% 的用户。