380. Insert Delete GetRandom O(1)

Problem¶

Design a set-like class supporting all in average \(O(1)\):

insert(val): Add val if absent.
remove(val): Delete val if present.
getRandom(): Return a uniformly random current element.

Example:

state: {}
insert(1) -> True    state: {1}
insert(2) -> True    state: {1,2}
remove(1) -> True    state: {2}
getRandom() -> 2

Key trick¶

Use two structures together:

A list of values for \(O(1)\) random access.
A dict value -> index for \(O(1)\) lookup/removal.

For deletion:

Swap the target with the last list element.
Update the moved element's index.
Pop the last element.

Trap¶

Common mistakes:

Removing from the middle of a list directly, which is \(O(n)\).
Using only a set, then converting to list in getRandom(), which is \(O(n)\).
Forgetting to update the index of the swapped last element.
Mishandling the case where the removed value is already the last element.

Why is it interesting?¶

It tests whether you can combine data structures so each one covers the other's weakness:

Dict gives fast membership/index.
List gives fast random access and tail deletion.

Python solution¶

import random


class RandomizedSet:
    def __init__(self):
        # values[i] = stored value at index i
        self.values = []

        # index[val] = where val is stored in values
        self.index = {}

    def insert(self, val: int) -> bool:
        # Already present
        if val in self.index:
            return False

        # Append to the end and record its index
        self.index[val] = len(self.values)
        self.values.append(val)
        return True

    def remove(self, val: int) -> bool:
        # Not present
        if val not in self.index:
            return False

        # Index of value to remove
        remove_idx = self.index[val]

        # Last value in the list
        last_val = self.values[-1]

        # Move last value into remove_idx so deletion stays O(1)
        self.values[remove_idx] = last_val
        self.index[last_val] = remove_idx

        # Remove last slot and delete mapping for val
        self.values.pop()
        del self.index[val]
        return True

    def getRandom(self) -> int:
        # random.choice on a list is O(1)
        return random.choice(self.values)

Complexity:

insert: \(O(1)\) average
remove: \(O(1)\) average
getRandom: \(O(1)\)

Comment on my solution¶

Your insert and remove are fine with a set, but getRandom breaks the requirement:

list(self.nums) rebuilds a full list every call.
That makes getRandom() \(O(n)\), not \(O(1)\).

So the missing idea is:

Keep a persistent list for random access.
Keep a dict for positions so removals stay \(O(1)\).

## Solution

import random

# Implemented after reading AI comments
# Wrong - because keeping removed elements in self.nums make these
#         elements wrong candidate when we call getRandom.
class RandomizedSet:

    def __init__(self):
        self.num_to_idx = {}
        self.nums = []


    def insert(self, val: int) -> bool:
        if val in self.num_to_idx:
            return False
        self.nums.append(val)
        self.num_to_idx[val] = len(self.nums) - 1
        return True

    def remove(self, val: int) -> bool:
        if val in self.num_to_idx:
            idx = self.num_to_idx[val]
            del self.num_to_idx[val]
            return True
        return False

    def getRandom(self) -> int:
        return self.nums[random.randrange(0, len(self.nums))]

# Works - but getRandom() is not O(1)
class RandomizedSet:

    def __init__(self):
        self.nums = set()


    def insert(self, val: int) -> bool:
        if val in self.nums:
            return False
        self.nums.add(val)
        return True

    def remove(self, val: int) -> bool:
        if val in self.nums:
            self.nums.remove(val)
            return True
        return False

    def getRandom(self) -> int:
        # I don't think this is O(1)
        return list(self.nums)[random.randrange(0, len(self.nums))]