- i1, j1, i2, j2 = self.rec_coo(A)
- 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
+ for n, X, f_X in [(1, A, f_A), (3, B, f_B)]:
+ for _ in range(torch.randint(n, (1,)) + 1):
+ i1, j1, i2, j2 = self.rec_coo(X)
+ X[i1:i2, j1:j2] = c1
+ f_X[i1:i2, j1:j2] = c2
+
+ def task_move(self, A, f_A, B, f_B):
+ c = torch.randperm(len(self.colors) - 1)[:1] + 1
+ di, dj = torch.randint(2, (2,)) * 2 - 1
+ for n, X, f_X in [(1, A, f_A), (3, B, f_B)]:
+ for _ in range(torch.randint(n, (1,)) + 1):
+ while True:
+ i1, j1, i2, j2 = self.rec_coo(X)
+ if (
+ i1 + di >= 0
+ and i2 + di < X.size(0)
+ and j1 + dj >= 0
+ and j2 + dj < X.size(1)
+ ):
+ break
+
+ X[i1:i2, j1:j2] = c
+ f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c
+
+ def task_grow(self, A, f_A, B, f_B):
+ c = torch.randperm(len(self.colors) - 1)[:1] + 1
+
+ for n, X, f_X in [(1, A, f_A), (3, B, f_B)]:
+ for _ in range(torch.randint(n, (1,)) + 1):
+ while True:
+ i1, j1, i2, j2 = self.rec_coo(X)
+ if i1 + 3 < i2 and j1 + 3 < j2:
+ break
+
+ X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c
+ f_X[i1:i2, j1:j2] = c