381. O(1) 时间插入、删除和获取随机元素 - 允许重复

困难设计力扣原题 ↗

在线写代码并运行看动画推演问 AI 讲解

题目描述

实现一个允许重复元素的多重集合类 RandomizedCollection：

RandomizedCollection()：创建空集合。insert(val)：无论是否已有，都加入一个 val；如果加入之前集合里没有 val，返回 true，否则返回 false。

remove(val)：如果集合里有 val，只删掉其中一个并返回 true；没有则返回 false。

getRandom()：随机返回集合中的一个元素。某个值被返回的概率与它在集合中的个数成正比。调用时集合保证非空。

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

示例 1

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

解释：插入后集合为 [1,1,2]，第一次随机返回了 1；删掉一个 1 后为 [1,2]，第二次随机返回了 2。随机结果每次运行可能不同。

数据范围

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

解题思路

在 380 的基础上，pos 改为“值 → 下标集合”。删除时从 pos[val] 取任意一个下标 i，把末尾元素 last 搬到 i：先把 i 加入 pos[last]，再从 pos[last] 删去末尾下标，最后 pop。

为什么正确

arr 紧凑，且对每个下标 i，i ∈ pos[arr[i]]。

复杂度分析

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

易错点

如果先删末尾下标再加 i，当 i 恰好等于末尾下标时会把一个已弹出的下标留在 pos 里，之后越界。

参考代码

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

class RandomizedCollection:
    def __init__(self):
        self.arr = []
        self.pos = defaultdict(set)

    def insert(self, val: int) -> bool:
        self.pos[val].add(len(self.arr))
        self.arr.append(val)
        return len(self.pos[val]) == 1

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

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

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

    public RandomizedCollection() {
    }

    public boolean insert(int val) {
        pos.computeIfAbsent(val, k -> new HashSet<>()).add(arr.size());
        arr.add(val);
        return pos.get(val).size() == 1;
    }

    public boolean remove(int val) {
        Set<Integer> s = pos.get(val);
        if (s == null || s.isEmpty())
            return false;
        int i = s.iterator().next(), last = arr.get(arr.size() - 1);
        s.remove(i);
        arr.set(i, last);
        pos.computeIfAbsent(last, k -> new HashSet<>()).add(i);
        pos.get(last).remove(arr.size() - 1);
        arr.remove(arr.size() - 1);
        return true;
    }

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

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

    bool insert(int val) {
        pos[val].insert(arr.size());
        arr.push_back(val);
        return pos[val].size() == 1;
    }

    bool remove(int val) {
        if (pos[val].empty())
            return false;
        int i = *pos[val].begin(), last = arr.back();
        pos[val].erase(i);
        arr[i] = last;
        pos[last].insert(i);
        pos[last].erase(arr.size() - 1);
        arr.pop_back();
        return true;
    }

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

type RandomizedCollection struct {
	arr []int
	pos map[int]map[int]bool
}

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

func (this *RandomizedCollection) Insert(val int) bool {
	if this.pos[val] == nil {
		this.pos[val] = map[int]bool{}
	}
	this.pos[val][len(this.arr)] = true
	this.arr = append(this.arr, val)
	return len(this.pos[val]) == 1
}

func (this *RandomizedCollection) Remove(val int) bool {
	if len(this.pos[val]) == 0 {
		return false
	}
	i := 0
	for k := range this.pos[val] {
		i = k
		break
	}
	delete(this.pos[val], i)
	last := this.arr[len(this.arr)-1]
	this.arr[i] = last
	if this.pos[last] == nil {
		this.pos[last] = map[int]bool{}
	}
	this.pos[last][i] = true
	delete(this.pos[last], len(this.arr)-1)
	this.arr = this.arr[:len(this.arr)-1]
	return true
}

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

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

RandomizedCollection.prototype.insert = function(val) {
    if (!this.pos.has(val)) this.pos.set(val, new Set());
    this.pos.get(val).add(this.arr.length);
    this.arr.push(val);
    return this.pos.get(val).size === 1;
};

RandomizedCollection.prototype.remove = function(val) {
    const s = this.pos.get(val);
    if (!s || s.size === 0)
        return false;
    const i = s.values().next().value, last = this.arr[this.arr.length - 1];
    s.delete(i);
    this.arr[i] = last;
    if (!this.pos.has(last)) this.pos.set(last, new Set());
    this.pos.get(last).add(i);
    this.pos.get(last).delete(this.arr.length - 1);
    this.arr.pop();
    return true;
};

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

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

    constructor() {
    }

    insert(val: number): boolean {
        if (!this.pos.has(val)) this.pos.set(val, new Set());
        this.pos.get(val)!.add(this.arr.length);
        this.arr.push(val);
        return this.pos.get(val)!.size === 1;
    }

    remove(val: number): boolean {
        const s = this.pos.get(val);
        if (!s || s.size === 0)
            return false;
        const i = s.values().next().value as number, last = this.arr[this.arr.length - 1];
        s.delete(i);
        this.arr[i] = last;
        if (!this.pos.has(last)) this.pos.set(last, new Set());
        this.pos.get(last)!.add(i);
        this.pos.get(last)!.delete(this.arr.length - 1);
        this.arr.pop();
        return true;
    }

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

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

设计专题的其他题目

查看设计全部题目 →