O(n) constructor, O(1) pick, two ways

  • 5

    Update: Second way

    I found another way. If a number only appears at one index, map the number to that index. Otherwise map it to a list of its indexes. Also got accepted all five times I submitted it.

    class Solution(object):
        def __init__(self, nums):
            indexes = self.indexes = {}
            for i, num in enumerate(nums):
                I = indexes.get(num)
                if I is None:
                    indexes[num] = i
                elif isinstance(I, int):
                    indexes[num] = [I, i]
        def pick(self, target):
            I = self.indexes[target]
            return I if isinstance(I, int) else random.choice(I)


    After quite a fight, this is my first version of this idea not getting "Memory Limit Exceeded". It got accepted all five times I submitted it.

    class Solution(object):
        def __init__(self, nums):
            count = {}
            for num in nums:
                count[num] = count.get(num, 0) + 1
            start, startstop = 0, count
            for num in count:
                startstop[num], start = (start << 32) | start, start + count[num]
            indexes = [None] * len(nums)
            for i, num in enumerate(nums):
                indexes[startstop[num] & 0xFFFFFFFF] = i
                startstop[num] += 1
            self.indexes = indexes
            self.startstop = startstop
        def pick(self, target):
            ss = self.startstop[target]
            return self.indexes[random.randrange(ss >> 32, ss & 0xFFFFFFFF)]

    First I count each number, then I write each number's indexes consecutively into indexes. So that for number num, its indexes are stored in indexes[start:stop]. And its start/stop are stored as startstop[num] = (start << 32) | stop.

    I tried this because the obvious solution storing one list of indexes for each number consistently got MLE, with *"Last executed input" having nums = [1,2,3,3,3]. I believe that that display is wrong, that I actually got MLE for a later case where nums has many more different numbers. Each different number having its own list can be expensive, so I came up with the above solution where I only have one list of all indexes, and each different number now has its own startstop int instead of its own list. An int takes only 24 bytes while even the empty list takes 72 bytes (as you can see with print sys.getsizeof([]) and print sys.getsizeof(2**60)).

    Note that I reuse the count dict as startstop instead of creating a new dict. That was the last optimization, the one that finally got me from MLE to Accepted.

  • 0

    The second way is brilliant!

  • 0

    It seems that these two methods are not feasible in Java.
    I coding the Original way in java, like below, but it still MLE

    public class Solution {
        private Random random;
        private int[] index;
        private Map<Integer, Long> startEnd;
        public Solution(int[] nums) {
            random = new Random();
            Map<Integer, Long> map = new HashMap();
            for (int i = 0; i < nums.length; i++) {
                Long res = map.get(nums[i]);
                res = res == null ? 1 : res + 1;
                map.put(nums[i], res);
            startEnd = map;
            Iterator<Integer> iterator = map.keySet().iterator();
            long start = 0;
            while (iterator.hasNext()) {
                int key = iterator.next();
                long value = map.get(key);
                startEnd.put(key, start << 32 | start);
                start += value;
            index = new int[nums.length];
            for (int i = 0; i < nums.length; i++) {
                Long value = startEnd.get(nums[i]);
                index[(int) (value & 0xFFFFFFFF)] = i;
                startEnd.put(nums[i], ++value);
        public int pick(int target) {
            long res = startEnd.get(target);
            int start = (int) (res >> 32);
            int end = (int) (res & 0xFFFFFFFF);
            return index[random.nextInt(end - start) + start];

  • 1

    @StefanPochmann Your second way gets me MLE.

  • 1

    @agave Hmm, I just submitted it five times and indeed it got MLE every time. Like I said, it did get accepted in five of five attempts before I posted it. I guess the test suite was changed or the limit lowered.

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.