Merge branch 'master' of hongabar.org:robotgame

robotgame/logic/colourboxes.py

@@ -27,11 +27,14 @@
 Colour boxes.
 """
 
+from __future__ import print_function
 import random
+import itertools
 
 def generate_colour_boxes(nwells, nboxes):
     """
-    Generate colour boxes that can be used to make all wells white.
+    Generate colour boxes that can be used to make all wells white. None of the
+    generated colour boxes are white.
 
     Arguments:
     nwells -- number of wells
@@ -44,18 +47,33 @@ def generate_colour_boxes(nwells, nboxes):
     nbits = nwells * 3
     data = [[0 for _ in range(nboxes)] for _ in range(nbits)]
 
-    def insert_1():
+    def insert_1(x):
         t = random.randrange(0, nboxes)
+        x0, x1 = _get_oxs(x)
         for y in range(t, nboxes) + range(0, t):
-            if data[x][y] == 0:
+            if data[x][y] == 0 and not (data[x0][y] == 1 and data[x1][y] == 1):
+                data[x][y] = 1
+                break
+        else:
+            raise Exception("Cannot maintain no 111s invariant.")
+
+    def insert_two_1(x):
+        t = random.randrange(0, nboxes)
+        x0, x1 = _get_oxs(x)
+        if len(list(filter(lambda y: data[x][y] == 0 and not (data[x0][y] == 1 and data[x1][y] == 1),
+                           range(nboxes)))) < 2:
+            return
+        for _ in range(2):
+            for y in range(t, nboxes) + range(0, t):
+                if data[x][y] == 0 and not (data[x0][y] == 1 and data[x1][y] == 1):
                     data[x][y] = 1
                     break
 
     for x in range(len(data)):
-        insert_1()
+        insert_1(x)
+    for x in range(len(data)):
         for _ in range(random.randrange(0, (nboxes + 1) / 2)):
-            insert_1()
-            insert_1()
+            insert_two_1(x)
 
     boxes = []
     for y in range(nboxes):
@@ -68,9 +86,19 @@ def generate_colour_boxes(nwells, nboxes):
             box.append((r, g, b))
     return boxes
 
+def _get_oxs(x):
+    return (x + 1, x + 2) if x % 3 == 0 \
+        else (x - 1, x + 1) if x % 3 == 1 \
+        else (x - 2, x - 1)
+
 def generate_random_box(nwells):
-    r = lambda: random.choice((0, 1))
-    return [(r(), r(), r()) for _ in range(nwells)]
+    """Generate a box with random colors except white (111)."""
+    def gen_wc():
+        wc = [random.choice((0, 1)) for i in range(3)]
+        if all(b == 1 for b in wc):
+            wc[random.randrange(3)] = 0
+        return wc
+    return [tuple(gen_wc()) for _ in range(nwells)]
 
 def get_colours(boxes):
     colours = []

tests/colourboxes_tests.py

@@ -6,8 +6,7 @@ from robotgame.logic.colourboxes import *
 
 class ColourboxesTest(unittest.TestCase):
     def test_generation(self):
-        for args in ((1, 1),
-                     (2, 2),
+        for args in ((1, 2),
                      (3, 3),
                      (4, 6)):
             boxes = generate_colour_boxes(*args)