- A[i1:i2, j1:j2] = self.name2color["red"]
- f_A[i1:i2, j1:j2] = self.name2color["green"]
- i1, j1, i2, j2 = self.rec_coo(B)
- B[i1:i2, j1:j2] = self.name2color["red"]
- f_B[i1:i2, j1:j2] = self.name2color["green"]
+ A[i1:i2, j1:j2] = c1
+ f_A[i1:i2, j1:j2] = c2
+ for _ in range(3):
+ i1, j1, i2, j2 = self.rec_coo(B)
+ B[i1:i2, j1:j2] = c1
+ f_B[i1:i2, j1:j2] = c2
+
+ def move_color(self, A, f_A, B, f_B):
+ c1, c2 = torch.randperm(len(self.colors) - 1)[:2] + 1
+
+ i1, j1, i2, j2 = self.rec_coo(A)
+ A[i1:i2, j1:j2] = c1
+ f_A[i1:i2, j1:j2] = c1
+
+ while True:
+ i1, j1, i2, j2 = self.rec_coo(A)
+ if i2 < self.height - 1:
+ break
+ A[i1:i2, j1:j2] = c2
+ f_A[i1:i2, j1:j2] = c2