leetcode-go/leetcode/medium/519.random-flip-matrix.go

type Solution struct {
	m    map[int]int
	size int
	rows int
	cols int
}

func Constructor(n_rows int, n_cols int) Solution {
	return Solution{
		m:    make(map[int]int),
		size: n_rows * n_cols,
		rows: n_rows,
		cols: n_cols,
	}
}

func (this *Solution) Flip() []int {
	idx := rand.Intn(this.size)
	i := idx
	this.size--
	if val, ok := this.m[idx]; ok {
		idx = val
	} // The trick is very important!
	if val, ok := this.m[this.size]; ok {
		this.m[i] = val
	} else {
		this.m[i] = this.size
	}
	return []int{idx / this.cols, idx % this.cols}
}

func (this *Solution) Reset() {
	this.m = make(map[int]int)
	this.size = this.rows * this.cols
}

/**
 * Your Solution object will be instantiated and called as such:
 * obj := Constructor(n_rows, n_cols);
 * param_1 := obj.Flip();
 * obj.Reset();
 */