Implemented three in a row but not the best way. Needs to be improved a lot

sources/Square.py

@@ -28,16 +28,24 @@ class Square:
         self.horiz = horiz
         self.vert = vert
 
-    def get_next_horiz(self):
+    def next_horiz(self, width):
+        if self.horiz == width - 1:
+            return None
         return Square(self.horiz + 1, self.vert)
 
-    def get_prev_horiz(self):
+    def prev_horiz(self):
+        if self.horiz == 0:
+            return None
         return Square(self.horiz - 1, self.vert)
 
-    def get_next_vert(self):
+    def next_vert(self, height):
+        if self.vert == height - 1:
+            return None
         return Square(self.horiz, self.vert + 1)
 
-    def get_prev_vert(self):
+    def prev_vert(self):
+        if self.vert == 0:
+            return None
         return Square(self.horiz, self.vert - 1)
 
     def __eq__(self, other):
@@ -51,12 +59,124 @@ class Square:
 
 
 def solve(grid):
-    for char in grid:
+    y = 0
-        pass
+    while y < len(grid):
+        x = 0
+        while x < len(grid[0]):
+            if grid[x][y] == ' ':
+                v_prev_x, v_prev_y = prev_vert(x, y)
+                v_p_prev_x, v_p_prev_y = prev_vert(v_prev_x, v_prev_y)
+                v_next_x, v_next_y = next_vert(x, y, len(grid))
+                v_n_next_x, v_n_next_y = next_vert(v_next_x, v_next_y,
+                                                   len(grid))
+
+                h_prev_x, h_prev_y = prev_horiz(x, y)
+                h_p_prev_x, h_p_prev_y = prev_horiz(h_prev_x, h_prev_y)
+                h_next_x, h_next_y = next_horiz(x, y, len(grid[0]))
+                h_n_next_x, h_n_next_y = next_horiz(h_next_x, h_next_y,
+                                                    len(grid[0]))
+
+                be_blue = (exists(v_prev_x, v_prev_y) and
+                           exists(v_p_prev_x, v_p_prev_y) and
+                           grid[v_prev_x][v_prev_y] == 'R' and
+                           grid[v_p_prev_x][v_p_prev_y] == 'R') or \
+                          (exists(v_next_x, v_next_y) and
+                           exists(v_n_next_x, v_n_next_y) and
+                           grid[v_next_x][v_next_y] == 'R' and
+                           grid[v_n_next_x][v_n_next_y] == 'R') or \
+                          (exists(h_prev_x, h_prev_y >= 0) and
+                           exists(h_p_prev_x, h_p_prev_y) and
+                           grid[h_prev_x][h_prev_y] == 'R' and
+                           grid[h_p_prev_x][h_p_prev_y] == 'R') or \
+                          (exists(h_next_x, h_next_y) and
+                           exists(h_n_next_y, h_n_next_y) and
+                           grid[h_next_x][h_next_y] == 'R' and
+                           grid[h_n_next_x][h_n_next_y] == 'R') or \
+                          (exists(h_prev_x, h_prev_y) and
+                           exists(h_next_x, h_next_y) and
+                           grid[h_prev_x][h_prev_y] == 'R' and
+                           grid[h_next_x][h_next_y] == 'R') or \
+                          (exists(v_prev_x, v_prev_y) and
+                           exists(v_next_x, v_next_y) and
+                           grid[v_prev_x][v_prev_y] == 'R' and
+                           grid[v_next_x][v_next_y] == 'R')
+
+                be_red = (v_prev_x >= 0 and v_prev_y >= 0 and
+                          v_p_prev_x >= 0 and v_p_prev_y >= 0 and
+                          grid[v_prev_x][v_prev_y] == 'B' and
+                          grid[v_p_prev_x][v_p_prev_y] == 'B') or \
+                         (v_next_x >= 0 and v_next_y >= 0 and
+                          v_n_next_x >= 0 and v_n_next_y >= 0 and
+                          grid[v_next_x][v_next_y] == 'B' and
+                          grid[v_n_next_x][v_n_next_y] == 'B') or \
+                         (h_prev_x >= 0 and h_prev_y >= 0 and
+                          h_p_prev_x >= 0 and h_p_prev_y >= 0 and
+                          grid[h_prev_x][h_prev_y] == 'B' and
+                          grid[h_p_prev_x][h_p_prev_y] == 'B') or \
+                         (h_next_x >= 0 and h_next_y >= 0 and
+                          h_n_next_y >= 0 and h_n_next_y >= 0 and
+                          grid[h_next_x][h_next_y] == 'B' and
+                          grid[h_n_next_x][h_n_next_y] == 'B') or \
+                          (exists(h_prev_x, h_prev_y) and
+                           exists(h_next_x, h_next_y) and
+                           grid[h_prev_x][h_prev_y] == 'B' and
+                           grid[h_next_x][h_next_y] == 'B') or \
+                          (exists(v_prev_x, v_prev_y) and
+                           exists(v_next_x, v_next_y) and
+                           grid[v_prev_x][v_prev_y] == 'B' and
+                           grid[v_next_x][v_next_y] == 'B')
+
+                if be_blue:
+                    grid[x][y] = 'B'
+                elif be_red:
+                    grid[x][y] = 'R'
+            x += 1
+        y += 1
 
 
-grid = [['R', ' ', ' ', 'R'],
+def exists(a, b):
-        ['R', ' ', 'B', ' '],
+    return a >= 0 and b >= 0
-        [' ', 'R', ' ', ' '],
-        [' ', ' ', ' ', ' ']]
 
+
+def next_horiz(horiz, vert, width):
+    if horiz == width - 1:
+        return -1, -1
+    return horiz + 1, vert
+
+
+def prev_horiz(horiz, vert):
+    if horiz == 0:
+        return -1, -1
+    return horiz - 1, vert
+
+
+def next_vert(horiz, vert, height):
+    if vert == height - 1:
+        return -1, -1
+    return horiz, vert + 1
+
+
+def prev_vert(horiz, vert):
+    if vert == 0:
+        return -1, -1
+    return horiz, vert - 1
+
+
+def print_grid(grid):
+    for line in grid:
+        for square in line:
+            if square == ' ':
+                square = '_'
+            print(square, end='  ')
+        print('')
+
+
+if __name__ == "__main__":
+    grid = [['B', ' ', ' ', ' '],
+            ['B', 'B', ' ', ' '],
+            [' ', ' ', ' ', ' '],
+            [' ', 'B', ' ', 'B']]
+    print_grid(grid)
+    solve(grid)
+    print("")
+    print_grid(grid)