Random Pick Index Python Solution with Tests

lcpy gen -n 398   # by problem number
lcpy gen -s random_pick_index   # by problem name

Input
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [1], [3]]
Output
[null, 4, 0, 2]

Explanation
Solution solution = new Solution([1, 2, 3, 3, 3]);
solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
solution.pick(1); // It should return 0. Since in the array only nums[0] is equal to 1.
solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.

import random
from collections import defaultdict


class Solution:
    # Time: O(n) init, O(1) pick
    # Space: O(n)
    def __init__(self, nums: list[int]) -> None:
        self.indices: dict[int, list[int]] = defaultdict(list)
        for i, num in enumerate(nums):
            self.indices[num].append(i)

    def pick(self, target: int) -> int:
        return random.choice(self.indices[target])