LeetCode 398. Random Pick Index

Description

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);
```

描述

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);

思路

• 这道题和第 382 题 Linked List Random Node 解法类似。
• 遍历给定的数组，统计当前满足条件的数的个数 n，用 random 函数随机生成一个 1 到 n 的数，如果选到了 n ，就把当前的数和最终结果替换。
```# -*- coding: utf-8 -*-
# @Author:             何睿
# @Create Date:        2019-08-30 19:55:46

import heapq
import random

from typing import List

class Solution:

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

def pick(self, target: int) -> int:
hep = []
for i, v in enumerate(self.arrays):
if v == target:
# heap 维护最小堆，random.random() 会被用来比较
# 没个数被弹出的概率都是一样的
heapq.heappush(hep, (random.random(), i))
_, index = heapq.heappop(hep)
return index

def pick2(self, target: int) -> int:

n, res = 0, 0
for x in filter(lambda x: x[1] == target, enumerate(self.arrays)):
n += 1
if random.randint(1, n) == n: # 解法思路同蓄水题解法思路
res = x[0]

return res
```

Posted in LeetCode and tagged , .

何 睿

This site uses Akismet to reduce spam. Learn how your comment data is processed.