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

• ## 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]
else:
indexes[num].append(i)

def pick(self, target):
I = self.indexes[target]
return I if isinstance(I, int) else random.choice(I)
``````

## Original:

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.

• The second way is brilliant!

• 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];
}
}
``````

• @StefanPochmann Your second way gets me MLE.

• @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.

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