diff options
Diffstat (limited to 'random_pick_index.py')
-rw-r--r-- | random_pick_index.py | 40 |
1 files changed, 40 insertions, 0 deletions
diff --git a/random_pick_index.py b/random_pick_index.py new file mode 100644 index 0000000..e79add9 --- /dev/null +++ b/random_pick_index.py @@ -0,0 +1,40 @@ +# 398. Random Pick Index +from random import randint + +""" +given an integer array 'nums' with possible duplicates, randomly output the +index of a given 'target' numbber. you can assume that the given target +number must exist in the array. implement the 'Solution' class: 'Solution(int +[] nums)' initiliases the object with the array 'nums', '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 +""" + + +class Solution(object): + def __init__(self, nums): + """ + :type nums: List[int] + """ + self.nums = nums + + def pick(self, target): + """ + :type target: int + :rtype: int + """ + cnt, res = 0, None + for i, n in enumerate(self.nums): + if n == target: + cnt += 1 + if (randint(0, cnt) % cnt) == 0: + res = i + return res + + +if __name__ == "__main__": + obj = Solution([1, 2, 3, 3, 3, 3]) + print(obj.pick(3)) + print(obj.pick(1)) + print(obj.pick(3)) |