380. O(1) 时间插入、删除和获取随机元素

题目描述

实现一个集合类 RandomizedSet，元素不重复：

RandomizedSet()：创建空集合。insert(val)：若 val 不在集合中，就加入并返回 true；已在集合中则不改动，返回 false。remove(val)：若 val 在集合中，就删除并返回 true；否则返回 false。

getRandom()：从当前集合中随机返回一个元素，每个元素被选中的概率必须相同。调用时集合保证非空。

每个方法的平均时间复杂度都必须是 O(1)。示例中第一行为方法名，第二行为参数，构造函数的返回记为 null。

示例 1

输入：操作 = ["RandomizedSet","insert","remove","insert","getRandom","remove","insert","getRandom"]
参数 = [[],[1],[2],[2],[],[1],[2],[]]
输出：[null,true,false,true,2,true,false,2]

解释：最后集合只剩 2，所以最后一次 getRandom() 必然返回 2；第一次随机调用时集合为 {1,2}，这里返回了 2。

数据范围

-2^31 <= val <= 2^31 - 1
insert、remove、getRandom 总共最多调用 2 * 10^5 次
调用 getRandom 时集合中至少有一个元素。

解题思路

数组 arr 保存元素，支持按下标随机访问；哈希表 pos 记录每个值的下标。删除时把末尾元素搬到被删位置、更新它的下标，再 pop 末尾，避免数组中间出现空洞。

为什么正确

arr 紧凑无空洞，且对每个元素 pos[arr[i]] == i。

复杂度分析

三个操作均为平均 O(1)；空间 O(n)。

易错点

删除时先更新 pos[last] 再删 pos[val]；如果当 val 就是末尾元素时先删后写，会把已删除的值重新写回哈希表。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class RandomizedSet:
    def __init__(self):
        self.arr = []
        self.pos = {}

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

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

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

class RandomizedSet {
    private final List<Integer> arr = new ArrayList<>();
    private final Map<Integer, Integer> pos = new HashMap<>();
    private final Random rnd = new Random();

    public RandomizedSet() {
    }

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

    public boolean remove(int val) {
        if (!pos.containsKey(val))
            return false;
        int i = pos.get(val), last = arr.get(arr.size() - 1);
        arr.set(i, last);
        pos.put(last, i);
        arr.remove(arr.size() - 1);
        pos.remove(val);
        return true;
    }

    public int getRandom() {
        return arr.get(rnd.nextInt(arr.size()));
    }
}

class RandomizedSet {
    vector<int> arr;
    unordered_map<int, int> pos;
public:
    RandomizedSet() {}

    bool insert(int val) {
        if (pos.count(val))
            return false;
        pos[val] = arr.size();
        arr.push_back(val);
        return true;
    }

    bool remove(int val) {
        if (!pos.count(val))
            return false;
        int i = pos[val], last = arr.back();
        arr[i] = last;
        pos[last] = i;
        arr.pop_back();
        pos.erase(val);
        return true;
    }

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

type RandomizedSet struct {
	arr []int
	pos map[int]int
}

func Constructor() RandomizedSet {
	return RandomizedSet{[]int{}, map[int]int{}}
}

func (this *RandomizedSet) Insert(val int) bool {
	if _, ok := this.pos[val]; ok {
		return false
	}
	this.pos[val] = len(this.arr)
	this.arr = append(this.arr, val)
	return true
}

func (this *RandomizedSet) Remove(val int) bool {
	i, ok := this.pos[val]
	if !ok {
		return false
	}
	last := this.arr[len(this.arr)-1]
	this.arr[i] = last
	this.pos[last] = i
	this.arr = this.arr[:len(this.arr)-1]
	delete(this.pos, val)
	return true
}

func (this *RandomizedSet) GetRandom() int {
	return this.arr[rand.Intn(len(this.arr))]
}

var RandomizedSet = function() {
    this.arr = [];
    this.pos = new Map();
};

RandomizedSet.prototype.insert = function(val) {
    if (this.pos.has(val))
        return false;
    this.pos.set(val, this.arr.length);
    this.arr.push(val);
    return true;
};

RandomizedSet.prototype.remove = function(val) {
    if (!this.pos.has(val))
        return false;
    const i = this.pos.get(val), last = this.arr[this.arr.length - 1];
    this.arr[i] = last;
    this.pos.set(last, i);
    this.arr.pop();
    this.pos.delete(val);
    return true;
};

RandomizedSet.prototype.getRandom = function() {
    return this.arr[Math.floor(Math.random() * this.arr.length)];
};

class RandomizedSet {
    private arr: number[] = [];
    private pos = new Map<number, number>();

    constructor() {
    }

    insert(val: number): boolean {
        if (this.pos.has(val))
            return false;
        this.pos.set(val, this.arr.length);
        this.arr.push(val);
        return true;
    }

    remove(val: number): boolean {
        if (!this.pos.has(val))
            return false;
        const i = this.pos.get(val)!, last = this.arr[this.arr.length - 1];
        this.arr[i] = last;
        this.pos.set(last, i);
        this.arr.pop();
        this.pos.delete(val);
        return true;
    }

    getRandom(): number {
        return this.arr[Math.floor(Math.random() * this.arr.length)];
    }
}

载入编辑器，自己跑一遍 →

设计专题的其他题目

查看设计全部题目 →