blueberries.py

from copy import deepcopy

# puzzle = '''\
#     ...12....
#     .3.......
#     22.......
#     .2.....4.
#     .....3.3.
#     .........
#     ..24...2.
#     .........
#     3......2.
# '''

puzzle = '''\
    .....1..1
    .........
    .......3.
    .........
    ..3...4..
    2..4.3...
    ..3....2.
    ......2..
    .3....0..
'''


puzzle = [[*line.strip()] for line in puzzle.strip().split("\n")]


class Blueberries:
    rows = [{(x, y) for x in range(9)} for y in range(9)]
    columns = [{(x, y) for y in range(9)} for x in range(9)]
    boxes = [{(b%3*3+i%3, b//3*3+i//3) for i in range(9)} for b in range(9)]
    regions = rows + columns + boxes

    def __init__(self, grid):
        self.start = deepcopy(grid)
        self.grid = deepcopy(grid)

    def neighbors(self, *args):
        x, y = args if len(args) == 2 else args[0]
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == dy == 0:
                    continue
                new_x = x + dx
                new_y = y + dy
                if 0 <= new_x < 9 and 0 <= new_y < 9:
                    yield (new_x, new_y)

    def count_around(self, x, y, alphabet):
        return sum(1 for (nx, ny) in self.neighbors(x, y) if self.grid[ny][nx] in alphabet)

    def count_in(self, region, alphabet):
        return sum(1 for (x, y) in region if self.grid[y][x] in alphabet)

    def is_forced_clue(self, x, y):
        return self.grid[y][x].isdigit() and \
            self.count_around(x, y, '*.') == int(self.grid[y][x])

    def is_saturated_clue(self, x, y):
        if not self.grid[y][x].isdigit():
            return False
        return self.count_around(x, y, '*') == int(self.grid[y][x]) and self.count_around(x, y, '.') > 0

    def clue_berry_step(self):
        changed = set()
        for y in range(9):
            for x in range(9):
                if self.is_forced_clue(x, y):
                    for nx, ny in self.neighbors(x, y):
                        if self.grid[ny][nx] == '.':
                            self.grid[ny][nx] = '*'
                            changed.add((x, y))
                if changed: return changed

    def clue_saturated_step(self):
        changed = set()
        for y in range(9):
            for x in range(9):
                if self.is_saturated_clue(x, y):
                    for nx, ny in self.neighbors(x, y):
                        if self.grid[ny][nx] == '.':
                            self.grid[ny][nx] = ' '
                            changed.add((x, y))
                if changed: return changed

    def region_berry_step(self):
        changed = set()
        for r in self.regions:
            if self.count_in(r, '.*') == 3:
                for (x, y) in r:
                    if self.grid[y][x] == '.':
                        self.grid[y][x] = '*'
                        changed |= r
                if changed: return changed

    def region_saturated_step(self):
        changed = set()
        for r in self.regions:
            if self.count_in(r, '*') == 3:
                for (x, y) in r:
                    if self.grid[y][x] == '.':
                        self.grid[y][x] = ' '
                        changed |= r
                if changed: return changed

    def basic_step(self):
        if clues := self.clue_berry_step():
            return ("All of \x1b[94mthis clue\x1b[0m's candidates must be berries.", clues)
        if clues := self.clue_saturated_step():
            return ("\x1b[94mThis clue\x1b[0m cannot see any more berries.", clues)
        if regions := self.region_saturated_step():
            return ("\x1b[94mThis region\x1b[0m can't have any more berries.", regions)
        if regions := self.region_berry_step():
            return ("All candidates in \x1b[94mthis region\x1b[0m must be berries.", regions)
        return None

    def contradiction(self):
        for r in self.regions:
            if self.count_in(r, '.*') < 3:
                return ("\x1b[94mThis region\x1b[0m couldn't contain three berries.", r)
            if self.count_in(r, '*') > 3:
                return ("\x1b[94mThis region\x1b[0m would have too many berries.", r)
        for y in range(9):
            for x in range(9):
                if self.grid[y][x].isdigit():
                    n = int(self.grid[y][x])
                    if self.count_around(x, y, '.*') < n:
                        return (f"\x1b[94mThis clue\x1b[0m could never see {n} berries.", {(x, y)})
                    if self.count_around(x, y, '*') > n:
                        return (f"\x1b[94mThis clue\x1b[0m would see too many berries.", {(x, y)})

    def find_contradiction(self, x, y, value):
        hypo = Blueberries(self.grid)
        hypo.grid[y][x] = value
        difficulty = 0
        while hypo.basic_step() and not hypo.contradiction():
            difficulty += 1
        c = hypo.contradiction()
        return (difficulty, c) if c else None

    def hypothetical_step(self):
        contradictions = []
        for y in range(9):
            for x in range(9):
                if self.grid[y][x] == '.':
                    if cy := self.find_contradiction(x, y, '*'):
                        rv = ("\x1b[91mThis cell\x1b[0m must be empty, as otherwise...\n" + cy[1][0], cy[1][1], {(x, y)})
                        contradictions.append((cy, x, y, ' ', rv))
                    if cn := self.find_contradiction(x, y, ' '):
                        rv = ("\x1b[91mThis cell\x1b[0m must be a berry, as otherwise...\n" + cn[1][0], cn[1][1], {(x, y)})
                        contradictions.append((cn, x, y, '*', rv))

        if contradictions:
            c, x, y, v, rv = min(contradictions)
            self.grid[y][x] = v
            return rv

    def advanced_step(self):
        return self.basic_step() or self.hypothetical_step()


    def __str__(self):
        return '\n'.join(map(' '.join, self.grid))

    def highlight(self, blue, red=()):
        def show(x, y):
            gbg = 254 if x//3%2^y//3%2 else 253
            char = self.grid[y][x].replace('.', '·').replace(' ', 'x')
            bg = 217 if (x,y) in red else 153 if (x,y) in blue else gbg
            sp = " " if x%3==0 else ""
            i = 1 if x%3<2 else 0
            rbg = 217 if (x+i,y) in red or (x,y) in red else \
                    153 if (x+i,y) in blue or (x,y) in blue else gbg
            fg = 160 if char=='x' else 33 if char=='*' else 16
            return f"\x1b[48;5;{bg}m{sp}\x1b[38;5;{fg}m{char}\x1b[48;5;{rbg}m \x1b[0m"
        def show_row(y):
            return f"    " + "".join(show(x, y) for x in range(9))

        return '\n'.join(show_row(y) for y in range(9))


puzzle = Blueberries(puzzle)
print("\x1b[2J\x1b[H" + "Blueberry Trio\n\n")
print(puzzle.highlight(()))
input()

while step := puzzle.advanced_step():
    print("\x1b[2J\x1b[H" + step[0] + "\n" * (2-step[0].count("\n")))
    print(puzzle.highlight(*step[1:]))
    input()

print("\x1b[2J\x1b[H" + "Puzzle solved!\n\n")
print(puzzle.highlight(()))
input()