$$ \newcommand\Tr{\mathrm{Tr}} \newcommand{\braket}[2]{\langle #1 \mid #2 \rangle} \newcommand\I{\mathbb{I}} \newcommand{\avg}[1]{\left< #1 \right>} \newcommand{\RD}{D} \newcommand{\ri}{\mathrm{i}} \DeclareMathOperator{\sign}{sign} \DeclareMathOperator{\Sign}{Sign} \newcommand{\ii}{\mathrm i} \newcommand{\vv}{\mathrm v} \newcommand{\ff}{\mathrm f} \newcommand{\mm}{\mathrm m} \newcommand{\ee}{\mathrm e} \newcommand{\xx}{\mathrm x} \newcommand{\RR}{\mathrm R} \newcommand{\dd}{\mathrm d} \newcommand{\FF}{\mathrm F} \newcommand{\BB}{\mathrm B} \newcommand{\vph}{v_{\mathrm{ph}}} $$

Insert Delete GetRandom O(1)

Problem

In this problem, our task is to design a data structure called RandomizedSet that allows us to perform Insert, Delete and GetRandom number in $O(1)$ time complexity on average.
Below is the specification:

  • RandomizedSet() will initialize RandomizedSet class.
  • bool insert(int num) will insert num into the set if not present and return true, false otherwise.
  • bool remove(int num) will remove num from the set if present and return true, false otherwise.
  • int getRandom() will retrieve the element of the current set with the same probability. It is guaranteed the set is not empty when this function is performed.

Solution

When we want to perform insert and delete in $O(1)$, then HashMap is a potential candidate. However, we will have a problem when we want to perform getRandom operation because if we rely on this data structure, then the map must be convereted to Array or List first that takes $O(N)$ before we randomize the index.
Therefore, we need to combine with other data structure possibly Array or List so we can avoid conversion process stated previously.
So, we will use below data structures:

  • HashMap where the key is the element and the value will be the index of the element in the ArrayList.
  • ArrayList to store the element contiguously that allows us to randomize the index and access the element in $O(1)$.

However, there is caveat for delete process.
Deletion in ArrayList generally has $O(N)$ time complexity. But, removing the last element of ArrayList has $O(1)$ time complexity.
Therefore, our algorithm should allow deletion only on the last element of ArrayList.

import kotlin.random.Random

class RandomizedSet() {
    private val map: MutableMap<Int, Int> = mutableMapOf()
    private val list: ArrayList<Int> = ArrayList()

    fun insert(num: Int): Boolean {
        if (map.containsKey(num))
            return false

        map[num] = list.size
        list.add(list.size, num)

        return true
    }

    fun remove(num: Int): Boolean {
        if (!map.containsKey(num))
            return false

        val lastElement = list[list.lastIndex]
        val indexToRemove = map[num]!!
        list[indexToRemove] = lastElement
        map[lastElement] = indexToRemove

        list.removeAt(list.lastIndex)
        map.remove(num)

        return true
    }

    fun getRandom(): Int {
        val index = Random.nextInt(list.size)
        return list[index]
    }

}

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Random;

public class RandomizedSet {
    private final HashMap<Integer, Integer> map;
    private final ArrayList<Integer> list;
    private final Random random;

    public RandomizedSet() {
        map = new HashMap<Integer, Integer>();
        list = new ArrayList<Integer>();
        random = new Random();
    }

    public boolean insert(int val) {
        if (map.containsKey(val))
            return false;

        map.put(val, list.size());
        list.add(list.size(), val);

        return true;
    }

    public boolean remove(int val) {
        if (!map.containsKey(val))
            return false;

        int lastElement = list.get(list.size() - 1);
        int indexToRemove = map.get(val);
        list.set(indexToRemove, lastElement);
        map.put(lastElement, indexToRemove);

        list.remove(list.size() - 1);
        map.remove(val);

        return true;
    }

    public int getRandom() {
        return list.get(random.nextInt(list.size()));
    }
}