Problem statement

https://leetcode.com/problems/random-flip-matrix/

Solution

Similar to problem 0380 Insert Delete GetRandom O(1) Random, but here we do not need to delete elements. First of all, we can look at our matrix as at just list. Next, we need to keep dictionary dic of flipped elements, for example we have [0, 1, 2, 3, 4, 5] places, and we choose element number 2, then we put dic[2] = 5 and decrease size of our set by one (virtual size). So next time we generate uniformly from [0, 1, 2, 3, 4] and if we again get 2, we return dic[2] = 5, so in fact we generate from [0, 1, dic[2], 3, 4]. Then, if we again get 2, we generate from [0, 1, dic[2], 3], where dic[2] = 4. Here we use .get(a, b) function, which returns value of key a if we found this key in dictionary and else return b.

Complexity

Time complexity for flip() will be O(1). Space complexity will be O(k), the maximum size of our dic.

Code

class Solution:
    def __init__(self, n_rows, n_cols):
        self.r, self.c = n_rows, n_cols
        self.size = self.r*self.c - 1
        self.dic = {}
        
    def flip(self):
        rand = random.randint(0, self.size)
        res = self.dic.get(rand, rand)
        self.dic[rand] = self.dic.get(self.size, self.size)
        self.size -= 1
        return divmod(res, self.c)
        
    def reset(self):
        self.size = self.r*self.c - 1
        self.dic = {}