O(1) Randomized Set and Array Shuffling

Insert Delete GetRandom O(1) (LeetCode 380)

Design a data structure that supports $O(1)$ expected time complexity for insert, remove, and getRandom.

Strategy: Combine a HashMap (for $O(1)$ lookups) and a Dynamic Array (ArrayList) (for $O(1)$ random access).

class RandomizedSet {
    private Map<Integer, Integer> map = new HashMap<>(); // val -> index in list
    private List<Integer> list = new ArrayList<>();
    private Random rand = new Random();

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

    public boolean remove(int val) {
        if (!map.containsKey(val)) return false;
        
        int idx = map.get(val);
        int lastElement = list.get(list.size() - 1);
        
        // Swap with last element to avoid O(N) shift
        list.set(idx, lastElement);
        map.put(lastElement, idx);
        
        // Remove the duplicated last reference
        list.remove(list.size() - 1);
        map.remove(val);
        return true;
    }

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

[!IMPORTANT] The "Swap with Last" trick is the cornerstone of $O(1)$ deletions in array-backed structures where order doesn't matter.

Allowed Duplicates (LeetCode 381)

Supporting multiple instances of the same value. The HashMap now stores a Set of indices for each value.

class RandomizedCollection {
    private Map<Integer, Set<Integer>> map = new HashMap<>(); // val -> set of indices
    private List<Integer> list = new ArrayList<>();
    private Random rand = new Random();

    public boolean insert(int val) {
        boolean absent = !map.containsKey(val) || map.get(val).isEmpty();
        map.computeIfAbsent(val, k -> new HashSet<>()).add(list.size());
        list.add(val);
        return absent;
    }

    public boolean remove(int val) {
        if (!map.containsKey(val) || map.get(val).isEmpty()) return false;
        
        int removeIdx = map.get(val).iterator().next();
        map.get(val).remove(removeIdx);
        
        int lastElement = list.get(list.size() - 1);
        // If not removing the last element, perform swap
        if (removeIdx != list.size() - 1) {
            list.set(removeIdx, lastElement);
            map.get(lastElement).remove(list.size() - 1);
            map.get(lastElement).add(removeIdx);
        }
        
        list.remove(list.size() - 1);
        return true;
    }

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

Shuffle an Array (LeetCode 384)

The Fisher-Yates Shuffle Algorithm ensures every permutation is equally likely (probability $1/N!$).

class Solution {
    private int[] original;
    private int[] array;
    private Random rand = new Random();

    public Solution(int[] nums) {
        original = nums.clone();
        array = nums.clone();
    }

    public int[] reset() {
        array = original.clone();
        return array;
    }

    public int[] shuffle() {
        for (int i = array.length - 1; i > 0; i--) {
            // Pick a random index from [0, i] and swap with i
            int j = rand.nextInt(i + 1);
            int temp = array[i];
            array[i] = array[j];
            array[j] = temp;
        }
        return array;
    }
}

Reservoir Sampling

How to pick $K$ random elements from a stream of unknown/infinite length with equal probability.

• For the $i$-th element in the stream:
• Generate a random number $r$ in $[0, i]$.
• If $r < K$, replace a random element in the reservoir of size $K$ with the $i$-th element.

// Logic for K=1: Equal probability 1/N
public int reservoirSampling(int[] stream) {
    int result = stream[0];
    Random rand = new Random();
    for (int i = 1; i < stream.length; i++) {
        // With probability 1/(i+1), update result
        if (rand.nextInt(i + 1) == 0) {
            result = stream[i];
        }
    }
    return result;
}

筋