Refactor group updates

strategy.py

@@ -122,67 +122,65 @@ class State:
                 assert row < s.nrows
                 i = row * COLUMNS + s.exa
                 if not s.locked(i):
-                    s.held = s.blocks[i]
-                    s.blocks[i] = NOBLOCK
                     s.colskip[s.exa] += 1
-                    s.ungroup(i)
+                    s.held = s.remove(i)
 
             elif move == DROP:
                 if s.held != NOBLOCK:
                     row = s.colskip[s.exa]
                     assert row > 0
                     i = (row - 1) * COLUMNS + s.exa
-                    s.blocks[i] = s.held
-                    s.held = NOBLOCK
                     s.colskip[s.exa] -= 1
-                    s.regroup(i)
+                    s.place(i, s.held)
+                    s.held = NOBLOCK
 
             elif move == SWAP:
                 assert not s.colbusy(s.exa)
                 row = s.colskip[s.exa]
                 i = row * COLUMNS + s.exa
                 j = i + COLUMNS
+
                 if j < len(s.blocks) and not s.locked(i) and not s.locked(j):
                     bi = s.blocks[i]
                     bj = s.blocks[j]
+                    assert bi != NOBLOCK
+                    assert bj != NOBLOCK
                     if bi != bj:
                         s.blocks[i] = NOBLOCK
                         s.blocks[j] = NOBLOCK
-                        s.ungroup(i)
-                        s.ungroup(j)
-                        s.blocks[j] = bi
-                        s.regroup(j)
-                        s.blocks[i] = bj
-                        s.regroup(i)
+                        visited = set()
+                        s.ungroup(i, visited)
+                        s.ungroup(j, visited)
+                        s.place(j, bi)
+                        s.place(i, bj)
 
         return s
 
-    def ungroup(self, i):
+    def remove(self, i):
+        block = self.blocks[i]
+        assert block != NOBLOCK
+        self.blocks[i] = NOBLOCK
+        self.ungroup(i, set())
+        return block
+
+    def ungroup(self, i, visited):
         assert self.blocks[i] == NOBLOCK
-        visited = set()
         oldid = self.groups[i]
 
         for nb in neighbors(i, self.blocks):
             if self.groups[nb] == oldid:
-                newgroup = self.scan_group(nb, visited)
-                if newgroup:
-                    self.maxgroup = newid = self.maxgroup + 1
-                    for j in newgroup:
-                        self.groups[j] = newid
-                        self.groupsizes[j] = len(newgroup)
+                self.create_group(nb, visited)
 
         self.groups[i] = 0
         self.groupsizes[i] = 0
 
-    def regroup(self, i):
-        assert self.blocks[i] != NOBLOCK
-        self.maxgroup = newid = self.maxgroup + 1
-        newgroup = self.scan_group(i, set())
-        for j in newgroup:
-            self.groups[j] = newid
-            self.groupsizes[j] = len(newgroup)
+    def place(self, i, block):
+        assert block != NOBLOCK
+        assert self.blocks[i] == NOBLOCK
+        self.blocks[i] = block
+        self.create_group(i, set())
 
-    def scan_group(self, start, visited):
+    def create_group(self, start, visited):
         def scan(i):
             if i not in visited:
                 yield i
@@ -192,7 +190,12 @@ class State:
                         yield from scan(nb)
 
         block = self.blocks[start]
-        return tuple(scan(start))
+        group = tuple(scan(start))
+        if group:
+            self.maxgroup = newid = self.maxgroup + 1
+            for j in group:
+                self.groups[j] = newid
+                self.groupsizes[j] = len(group)
 
     def fragmentation(self):
         """