Insert Delete GetRandom O(1)

mediumAmazon-tagged: Design-style coding32 stepsJava · Python · C++

Support insert, remove and getRandom in average O(1). Keep values in an ArrayList (for random access) and a map from value to its index. To remove, move the last element into the removed slot and pop the end, so the list never has holes.

▶ Watch the animated trace

Traced input: insert(1), remove(2), insert(2), getRandom(), remove(1), insert(2), getRandom() → [null, true, false, true, 2, true, false, 2]. Original problem: neetcode.io/problems/insert-delete-getrandom-o1

Complexity

BestO(1)

AverageO(1)

WorstO(1)

SpaceO(n)

Average O(1) because of hashing.

Key idea

An array gives O(1) random picks; a map gives O(1) lookups.
Removing from the middle is O(n), so swap with the last element and remove the end instead.
Update the moved element's index in the map.

Common mistakes

Forgetting to update the moved element's index.
Removing by index with list.remove(int) from the middle is O(n).
Using a HashSet alone: there's no O(1) random pick.

Solution

Java

class RandomizedSet {
    private final List<Integer> list = new ArrayList<>();
    private final Map<Integer, Integer> index = new HashMap<>();
    private final Random rand = new Random(42);

    boolean insert(int val) {
        if (index.containsKey(val)) {
            return false;
        }
        index.put(val, list.size());
        list.add(val);
        return true;
    }

    boolean remove(int val) {
        Integer i = index.get(val);
        if (i == null) {
            return false;
        }
        int last = list.get(list.size() - 1);
        list.set(i, last);
        index.put(last, i);
        list.remove(list.size() - 1);
        index.remove(val);
        return true;
    }

    int getRandom() {
        int pick = list.get(rand.nextInt(list.size()));
        return pick;
    }
}

Python

class RandomizedSet:
    def __init__(self):
        self.items = []
        self.index = {}

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

    def remove(self, val: int) -> bool:
        if val not in self.index:
            return False
        i, last = self.index[val], self.items[-1]
        self.items[i] = last
        self.index[last] = i
        self.items.pop()
        del self.index[val]
        return True

    def getRandom(self) -> int:
        return random.choice(self.items)

C++

class RandomizedSet {
    vector<int> items;
    unordered_map<int, int> index;

public:
    bool insert(int val) {
        if (index.count(val)) {
            return false;
        }
        index[val] = items.size();
        items.push_back(val);
        return true;
    }

    bool remove(int val) {
        auto it = index.find(val);
        if (it == index.end()) {
            return false;
        }
        int i = it->second, last = items.back();
        items[i] = last;
        index[last] = i;
        items.pop_back();
        index.erase(val);
        return true;
    }

    int getRandom() {
        return items[rand() % items.size()];
    }
};

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
Remove the last element The swap is with itself.	`insert(5), insert(6), remove(6)`	`[null, true, true, true]`
Duplicate insert Already present.	`insert(3), insert(3)`	`[null, true, false]`
Remove missing Nothing to remove.	`remove(9)`	`[null, false]`