Skip to content

Instantly share code, notes, and snippets.

@itarato

itarato/4_in_a_row.py

Last active September 16, 2025 00:48
Show Gist options
  • Save itarato/2d370b038e3b0e31347021a15a84670c to your computer and use it in GitHub Desktop.
Save itarato/2d370b038e3b0e31347021a15a84670c to your computer and use it in GitHub Desktop.
Download ZIP
4 in a row on terminal
Raw
4_in_a_row.py
import math
class FourInARow:
def __init__(self) -> None:
self.map = []
for _ in range(6):
self.map.append([EMPTY] * 7)
def calculate_column_scores(self) -> list[int]:
pass
def step(self, player: int, column: int) -> bool:
for y in range(6):
if self.map[y][column] == EMPTY:
self.map[y][column] = player
return True
return False
def unstep(self, column: int):
for y in range(5, -1, -1):
if self.map[y][column] != EMPTY:
self.map[y][column] = EMPTY
return
raise "Cannot unstep empty column"
def score(self) -> dict[int, float]:
out = {
PLAYER_HUMAN: 0,
PLAYER_CPU: 0,
}
for y in range(6):
for x in range(7):
if self.map[y][x] == EMPTY:
continue
out[self.map[y][x]] += self.cell_score(x, y)
return out
def cell_score(self, x: int, y: int) -> float:
player = self.map[y][x]
_score = 0
# Vertical: |
current_len = self.len_of(x, y, dx=0, dy=-1, match=[player])
grow_len = self.len_of(x, y + 1, dx=0, dy=1, match=[EMPTY])
_score += FourInARow.score_for_current_and_grow_len(current_len, grow_len)
# Left horizontal: <-
current_len = self.len_of(x, y, dx=1, dy=0, match=[player])
grow_len = self.len_of(
x - 1, y, dx=-1, dy=0, match=[EMPTY, player], firstMatch=[EMPTY]
)
_score += FourInARow.score_for_current_and_grow_len(current_len, grow_len)
# Right horizontal: ->
current_len = self.len_of(x, y, dx=-1, dy=0, match=[player])
grow_len = self.len_of(
x + 1, y, dx=1, dy=0, match=[EMPTY, player], firstMatch=[EMPTY]
)
_score += FourInARow.score_for_current_and_grow_len(current_len, grow_len)
# Left-diag: \
current_len = self.len_of(x, y, dx=1, dy=-1, match=[player])
grow_len = self.len_of(
x - 1, y + 1, dx=-1, dy=1, match=[EMPTY, player], firstMatch=[EMPTY]
)
_score += FourInARow.score_for_current_and_grow_len(current_len, grow_len)
# Right-diag: /
current_len = self.len_of(x, y, dx=-1, dy=-1, match=[player])
grow_len = self.len_of(
x + 1, y + 1, dx=1, dy=1, match=[EMPTY, player], firstMatch=[EMPTY]
)
_score += FourInARow.score_for_current_and_grow_len(current_len, grow_len)
return _score
def score_for_current_and_grow_len(current_len: int, grow_len: int) -> float:
if current_len >= 4:
# It's twice the winning score so we can safely ask if the DIFF (CPU-HUMAN) is greater than winning.
return WIN_SCORE * 2
elif current_len + grow_len >= 4:
return current_len**2
else:
return 0
def coord_is_in_range(x: int, y: int) -> bool:
return x >= 0 and y >= 0 and x <= 6 and y <= 5
def len_of(
self,
x: int,
y: int,
dx: int,
dy: int,
match: list[int],
firstMatch: list[int] = None,
) -> int:
out = 0
isFirst = True
if firstMatch == None:
firstMatch = match
while FourInARow.coord_is_in_range(x, y):
if isFirst:
isFirst = False
didMatch = self.map[y][x] in firstMatch
else:
didMatch = self.map[y][x] in match
if didMatch:
out += 1
else:
break
x += dx
y += dy
return out
def dump_board(self) -> None:
print("\x1b[91m●\x1b[0m: Human")
print("\x1b[93m●\x1b[0m: Computer")
for y in range(5, -1, -1):
print("\x1b[100m\x1b[91m \x1b[0m", end="")
for x in range(7):
c = "\x1b[100m \x1b[0m"
if self.map[y][x] == PLAYER_HUMAN:
c = "\x1b[100m\x1b[91m● \x1b[0m"
elif self.map[y][x] == PLAYER_CPU:
c = "\x1b[100m\x1b[93m● \x1b[0m"
print(c, end="")
print("")
print("\x1b[44m 0 1 2 3 4 5 6 \x1b[0m\n\n")
class BaseGame:
def __init__(self) -> None:
self.fiar = FourInARow()
def run(self):
isHumanStep = True
while True:
self.fiar.dump_board()
current_score = self.fiar.score()
if current_score[PLAYER_HUMAN] >= WIN_SCORE:
print("HUMAN WON!!!")
break
elif current_score[PLAYER_CPU] >= WIN_SCORE:
print("CPU WON!!!")
break
if isHumanStep:
if not self.human_step():
break
else:
self.cpu_step()
isHumanStep = not isHumanStep
def human_step(self) -> bool:
cmd = input("?> ")
if cmd == "q":
return False
else:
col = int(cmd)
if col < 0 or col >= 7:
print("Bad command")
return self.human_step()
self.fiar.step(PLAYER_HUMAN, col)
return True
def cpu_step(self):
raise "Must be implemented"
class SimpleScoreGame(BaseGame):
def cpu_step(self):
best_col = None
best_score = -1
for x in range(7):
if not self.fiar.step(PLAYER_CPU, x):
continue
score = self.fiar.score()
if score[PLAYER_CPU] > best_score:
best_score = score[PLAYER_CPU]
best_col = x
self.fiar.unstep(x)
if best_col == None:
raise "No more steps"
else:
self.fiar.step(PLAYER_CPU, best_col)
class MinimaxGame(BaseGame):
def __init__(self, depth: int = 4):
super().__init__()
self.dotfile = Dot()
self.dotfile.record_header()
self.depth = depth
def __del__(self):
self.dotfile.record_footer()
if hasattr(self, "dotfile") and self.dotfile and hasattr(self.dotfile, "file"):
try:
self.dotfile.file.close()
except Exception:
pass
def cpu_step(self):
_, col = self.minimax(True, self.depth)
self.fiar.step(PLAYER_CPU, col)
def minimax(
self,
is_cpu: bool,
depth: int,
dot_node_name: str = "root",
parent_bounds: tuple[float, float] = (-math.inf, math.inf),
) -> list:
score = self.fiar.score()
cpu_heuristic_val = score[PLAYER_CPU] - score[PLAYER_HUMAN]
node_bounds = (-math.inf, math.inf)
if (
depth <= 0
or abs(score[PLAYER_CPU]) >= WIN_SCORE
or abs(score[PLAYER_HUMAN]) >= WIN_SCORE
):
if (
abs(score[PLAYER_CPU]) >= WIN_SCORE
or abs(score[PLAYER_HUMAN]) >= WIN_SCORE
):
cpu_heuristic_val *= depth + 1
self.dotfile.record_node(
dot_node_name,
f"{cpu_heuristic_val}\nC{score[PLAYER_CPU]}\nH{score[PLAYER_HUMAN]}",
"blue" if is_cpu else "red",
)
return [cpu_heuristic_val, None]
elif is_cpu:
value = -math.inf
col = None
for x in range(7):
# Alpha/beta pruning
if ALPHA_BETA_PRUNING and node_bounds[0] > parent_bounds[1]:
continue
if not self.fiar.step(PLAYER_CPU, x):
continue
dot_subnode_name = f"{dot_node_name}::{depth}_{x}"
# subvalue (MUST BE) (at least) [value]
subvalue, _ = self.minimax(
not is_cpu, depth - 1, dot_subnode_name, parent_bounds=node_bounds
)
self.dotfile.record_edge(dot_node_name, dot_subnode_name)
if subvalue >= value:
value = subvalue
col = x
node_bounds = (subvalue, node_bounds[1])
self.fiar.unstep(x)
self.dotfile.record_node(
dot_node_name,
str(value) + "(" + str(col) + ")",
"blue",
)
return [value, col]
else:
value = math.inf
for x in range(7):
# Alpha/beta pruning
if ALPHA_BETA_PRUNING and node_bounds[1] < parent_bounds[0]:
continue
if not self.fiar.step(PLAYER_HUMAN, x):
continue
dot_subnode_name = f"{dot_node_name}::{depth}_{x}"
# subvalue (MUST BE) (maximum) [value]
subvalue, _ = self.minimax(
not is_cpu, depth - 1, dot_subnode_name, parent_bounds=node_bounds
)
self.dotfile.record_edge(dot_node_name, dot_subnode_name)
if subvalue <= value:
value = subvalue
col = x
node_bounds = (node_bounds[0], subvalue)
self.fiar.unstep(x)
self.dotfile.record_node(
dot_node_name,
str(value) + "(" + str(col) + ")",
"red",
)
return [value, col]
class NullDot:
def record_header(self):
pass
def record_footer(self):
pass
def record_node(self, name: str, label: any, color: str):
pass
def record_edge(self, lhs, rhs):
pass
class Dot:
def __init__(self):
self.file = open("/tmp/minimax.dot", "w")
self.file.truncate(0)
def record_header(self):
self.file.write("digraph Tree {\n")
def record_footer(self):
self.file.write("}\n")
def record_node(self, name: str, label: any, color: str):
self.file.write(f'"{name}" [label="{label}", color={color}];\n')
def record_edge(self, lhs, rhs):
self.file.write(f'"{lhs}" -> "{rhs}";\n')
def test_scenario_not_taking_win_step():
mmg = MinimaxGame(2)
mmg.fiar.step(PLAYER_HUMAN, 3)
mmg.fiar.step(PLAYER_HUMAN, 4)
mmg.fiar.step(PLAYER_CPU, 3)
mmg.fiar.dump_board()
print(mmg.fiar.score())
mmg.cpu_step()
mmg.fiar.dump_board()
EMPTY = None
PLAYER_HUMAN = 1
PLAYER_CPU = 2
WIN_SCORE = 1000
ALPHA_BETA_PRUNING = True
if __name__ == "__main__":
MinimaxGame(5).run()
# test_scenario_not_taking_win_step()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment