O(1) 時間插入、刪除和獲取隨機元素,允許元素重複

低吟不作語發表於2020-11-06

設計一個支援在平均時間複雜度 O(1) 下, 執行以下操作的資料結構。注意:允許出現重複元素。

  • insert(val) :向集合中插入元素 val
  • remove(val) :當 val 存在時,從集合中移除一個 val
  • getRandom() :從現有集合中隨機獲取一個元素。每個元素被返回的概率應該與其在集合中的數量呈線性相關

示例

// 初始化一個空的集合。
RandomizedCollection collection = new RandomizedCollection();

// 向集合中插入 1 。返回 true 表示集合不包含 1 。
collection.insert(1);

// 向集合中插入另一個 1 。返回 false 表示集合包含 1 。集合現在包含 [1,1] 。
collection.insert(1);

// 向集合中插入 2 ,返回 true 。集合現在包含 [1,1,2] 。
collection.insert(2);

// getRandom 應當有 2/3 的概率返回 1 ,1/3 的概率返回 2 。
collection.getRandom();

// 從集合中刪除 1 ,返回 true 。集合現在包含 [1,2] 。
collection.remove(1);

// getRandom 應有相同概率返回 1 和 2 。
collection.getRandom();

解題思路

用 List 存放元素,再用 Map 來記錄每個元素在 List 中的所有下標,按照這個思想,實現插入和獲取隨機數都不難,唯獨刪除操作會卡殼。每次刪除時不需要真的從 List 中將元素刪掉,只要用 List 中的最後一個元素覆蓋待刪除元素,在 Map 中修改最後一個元素的下標即可。需要注意的是,每次刪除操作後,List 的後段會有無效的部分,因為每次執行刪除操作後,有效部分的最後一個元素就變成無效了,所以要多加註意

class RandomizedCollection {

    // 存放元素
    private List<Integer> list;
    // 記錄每個元素在 list 中的所有下標
    private Map<Integer, List<Integer>> map;
    // 執行過的有效 remove 數
    private int removeCount;

    /** Initialize your data structure here. */
    public RandomizedCollection() {
        list = new ArrayList<>();
        map = new HashMap<>();
        removeCount = 0;
    }

    public boolean insert(int val) {
        boolean flag;
        // removeCount == 0,直接在 list 後追加元素
        if(removeCount == 0) {
            list.add(val);        
        } else {
            // 否則要在 list 的有效範圍內追加元素,由刪除操作思想可知,list 末尾的部分元素是無效的
            list.add(list.size() - removeCount, val);
        }
        // 是否已有重複元素存在
        if(map.containsKey(val)) {
            flag = false;
            map.get(val).add(list.size() - 1);
        } else {
            flag = true;
            List<Integer> indexList = new LinkedList<>();
            indexList.add(list.size() - 1);
            map.put(val, indexList);
        }
        return flag;
    }
    
    public boolean remove(int val) {
        boolean flag;
        // 元素存在並且對應的下標記錄數不為 0
        if(map.containsKey(val) && map.get(val).size() != 0) {
            flag = true;
            // 取出第一個元素所在的下標並用 index 儲存
            int index = map.get(val).get(0);
            // 從 indexList 中刪除第一個下標
            map.get(val).remove(0);
            // 獲取最後一個有效元素
            int lastVal = list.get(list.size() - 1 - removeCount);
            // lastVal 覆蓋 index 處的值(從 list 中刪除一個 val)
            list.set(index, lastVal);
            // 獲取 lastVal 所在的所有下標
            List<Integer> tempList = map.get(lastVal);
            // 遍歷 tempList,找到原本所處的 list 有效範圍末尾的下標,用 index 替代
            for(int i = 0; i < tempList.size(); i++) {
                if(tempList.get(i) == (list.size() - 1 - removeCount)) {
                    map.get(lastVal).set(i, index);
                    break;
                }
            }
            removeCount++;
        } else {
            flag = false;
        }
        return flag;
    }
    
    public int getRandom() {
        int randomNum = new Random().nextInt(list.size() - removeCount);
        return list.get(randomNum);
    }
}

相關文章