leetcode-go/leetcode/hard/381.insert-delete-getrandom-o1-duplicates-allowed.go

type pair struct {
	_1 int
	_2 *list.Element
}

type RandomizedCollection struct {
	nums []pair
	size int
	pos  map[int]*list.List // Hash map + linked list, O(1) get random, O(1) insert, O(1) remove (average)
}

/** Initialize your data structure here. */
func Constructor() RandomizedCollection {
	return RandomizedCollection{pos: make(map[int]*list.List)}
}

/** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
func (this *RandomizedCollection) Insert(val int) bool {
	li, ok := this.pos[val]
	if !ok {
		li = list.New()
		this.pos[val] = li
	}
	ele := li.PushBack(this.size)
	this.nums = append(this.nums, pair{val, ele})
	this.size++
	return !ok
}

/** Removes a value from the collection. Returns true if the collection contained the specified element. */
func (this *RandomizedCollection) Remove(val int) bool {
	li, ok := this.pos[val]
	if !ok {
		return false
	}
	idx := li.Remove(li.Front()).(int)
	this.size--
	this.nums[idx] = this.nums[this.size]
	this.nums[idx]._2.Value = idx
	this.nums = this.nums[:this.size]
	if li.Len() == 0 {
		delete(this.pos, val)
	}
	return true
}

/** Get a random element from the collection. */
func (this *RandomizedCollection) GetRandom() int {
	return this.nums[rand.Intn(this.size)]._1
}

/**
 * Your RandomizedCollection object will be instantiated and called as such:
 * obj := Constructor();
 * param_1 := obj.Insert(val);
 * param_2 := obj.Remove(val);
 * param_3 := obj.GetRandom();
 */