Knights and Pages puzzle from The Chicken from Minsk

knightspages.py

# This implements the "Knights and Pages" problem
# from The Chicken from Minsk, Chapter 1
#
# PROBLEM:
#   Many years ago three knights waited to cross the river Neva.
#   Each knight had his own page, so there were six people. The
#   boat they had could only carry two. However, the knights were
#   ferocious killers and the pages were terrified. In fact it was
#   certain that any one of the pages would die of heart failure
#   if he were not protected at every instant from the other knights
#   by the presence of his own master. Was there any way to cross the
#   river without losing a page?
#
# The rules are meant to be simple and obvious and so should be interpreted
# as such. Basically a page dies if it is in transit on the boat with a
# different knight OR if it is on either shore in the presence of other
# knights AND without its own to protect it.
#
# This code implements that puzzle plus some variants given as follow-ups.
#
import json
from dataclasses import dataclass
from itertools import combinations, repeat


# making classes for these is overkill but it's an excuse to learn "dataclass"
@dataclass(frozen=True)
class Being:
    id: int
    typechar: str

    def __str__(self):
        return f"{self.typechar}{self.id}"


@dataclass(frozen=True)
class Page(Being):
    typechar: str = 'P'


@dataclass(frozen=True)
class Knight(Being):
    typechar: str = 'K'


class Puzzle:
    # The boat is represented as location 0: self.locations[self.BOAT]
    BOAT = 0

    def __init__(self, *, kps=3, boatcapacity=2, nlocations=2):
        self.boatcapacity = boatcapacity
        self.locations = [set() for _ in range(nlocations + 1)]  # +1 for boat
        self.allkp = set()

        # this tracks the history of moves which is convenient for solver
        self.movetrail = []

        # everyone starts in location 1... (remember: 0 is the BOAT)
        self.boatlocation = 1
        self._forceonto(1,
                        *(Knight(i) for i in range(kps)),
                        *(Page(i) for i in range(kps)))

    def _forceonto(self, locn, *kps):
        """Just put knights/pages in locn. Adds to allkp too."""
        where = self.locations[locn]
        for kp in kps:
            where.add(kp)
            self.allkp.add(kp)

    def tojson(self):
        """Return JSON encoding."""
        d = {'boatcapacity': self.boatcapacity,
             'boatlocation': self.boatlocation,
             'movetrail': self.movetrail,
             'lists': [[str(x) for x in loc] for loc in self.locations]}
        return json.dumps(d)

    @classmethod
    def fromjson(cls, s):
        """Factory. Returns a new Puzzle from a tojson() string."""
        d = json.loads(s)

        # make one with no knights or pages
        z = cls(kps=0, boatcapacity=d['boatcapacity'],
                nlocations=len(d['lists']) - 1)     # -1 for boat

        for list_i, kps in enumerate(d['lists']):
            for x in kps:
                match x[0], int(x[1:]):
                    case ('K', id):
                        z._forceonto(list_i, Knight(id))
                    case ('P', id):
                        z._forceonto(list_i, Page(id))
                    case _:
                        raise ValueError(f"cannot decode {s}")

        z.boatlocation = d['boatlocation']
        z.movetrail = d['movetrail']
        z.raise_violations()
        return z

    def clone(self):
        """Factory. Return a new Puzzle with same state."""
        return self.fromjson(self.tojson())

    def knights(self, *, id=None, where=None):
        """Always returns a LIST.

        With no arguments: Returns LIST of all knights.
        With id and/or where: Returns LIST Knights that match.
        """
        return list(self._find(where, id, Knight))

    def pages(self, *, id=None, where=None):
        """Always returns a LIST.

        With no arguments: Returns LIST of all Pages.
        With id and/or where: Returns LIST Pages that match.
        """
        return list(self._find(where, id, Page))

    # convenience - find a specific Knight; raises IndexError if not found
    def knight_by_id(self, id, where=None):
        return self.knights(id=id, where=where)[0]

    # convenience - find a specific Page; raises IndexError if not found
    def page_by_id(self, id, where=None):
        return self.pages(id=id, where=where)[0]

    # general search function used by the above knights/pages/etc finders
    def _find(self, where, id, t):
        """Yield knights/pages matching criteria.

        If where is None, search everywhere, else only search within where
        If id is not None, yield only matching ids
        If t is not None, yield only matching types (Knight or Page)
        If t is tuple, yield only tuples of paired Knight/Page
        """

        if where is None:
            where = self.allkp
        for b in where:
            if t == tuple and isinstance(b, Knight):
                try:
                    yield b, self.page_by_id(b.id, where=where)
                except IndexError:
                    pass
            elif t is None or isinstance(b, t):
                if id is None or b.id == id:
                    yield b

    def violations(self, *x, reason="constraint violation: "):
        """RETURN exception if self violates a constraint; else None.

        With no arguments, enforces constraints on the entire puzzle state.
        With arguments, only checks the given lists.

        The rule is if any Page p is not accompanied by its own Knight,
        AND is in the presence of any other Knight, it dies.
        """

        for loc in x or self.locations:
            knights = set(k.id for k in self.knights(where=loc))
            if knights:
                pages = set(p.id for p in self.pages(where=loc))
                if pages - knights:
                    return ValueError(
                        f"{reason}" +
                        f"Page(s) '{pages-knights}' unprotected in '{loc}'")
        return None

    def raise_violations(self, *x, reason="constraint violation: "):
        """See find_violations; this raises any returned exception."""
        if exc := self.violations(*x, reason=reason):
            raise exc

    @property
    def endstate(self):
        return len(self.locations[-1]) == len(self.allkp)

    def _statevec(self, loc):
        """Return a 4-tuple representing state of given location.

        Result is: (kp, k, p, b)
        where
        kp:   knight/page pair count
        k:    unaccompanied knight count
        p:    unaccompanied page count
        b:    boat is here (True/False)
        """

        kp = k = p = 0
        knightIDs = set(x.id for x in self.knights(where=loc))
        pageIDs = set(x.id for x in self.pages(where=loc))
        for x in loc:
            if x.id not in knightIDs or x.id not in pageIDs:
                if isinstance(x, Knight):
                    k += 1
                else:
                    p += 1
            else:
                kp += 1
        return (kp, k, p, self.locations[self.boatlocation] is loc)

    def canonicalstate(self):
        """Return a representation of canonical puzzle state.

        The canonical representation ignores superficial numbering
        differences between otherwise-identical puzzle states.
        For example, these two puzzle states are the same situation:

            1)   (K0P0,K1P1)      (K2P2 + the boat)
            2)   (K0P0,K2P2)      (K1P1 + the boat)

        and will have identical canonical states.
        """
        # any violation is canonically identical
        if self.violations():
            return "VIOLATION"

        # Make all the states except for the last one, then sort them.
        # Sorting the states collapses more logically-identical cases
        # when there are intermediate islands (nlocations > 2).
        # For example:
        #                KP,B   P       K
        #          and   P      KP,B    K
        #
        # are also the same canonical state.
        #
        # NOTE: The last location must not be sorted as it is not fungible.
        #       As a trivial example of why, the initial state and the end
        #       state would have identical canonical states if the last
        #       location is treated identically to all other locations.
        #       Therefore, the last location is forced to stay last but the
        #       remaining locations are treated as fungible.

        states = sorted([self._statevec(loc) for loc in self.locations[:-1]])
        states.append(self._statevec(self.locations[-1]))
        return tuple(states)

    def all_moves(self):
        """Generate all moves, whether they are violations or not."""
        locs = set(range(len(self.locations)))
        locs.remove(self.boatlocation)
        locs.remove(self.BOAT)
        for n in range(self.boatcapacity):
            for loc in locs:
                yield from zip(
                    repeat(loc),
                    combinations(self.locations[self.boatlocation], n+1))

    def unique_legal_moves(self):
        """Generate legal moves without any logically identical dups."""
        canons = set()
        for loc, m in self.all_moves():
            z = self.clone()
            try:
                z.transfer(m, loc)
            except ValueError:
                pass
            else:
                c = z.canonicalstate()
                if c not in canons:
                    canons.add(c)
                    yield loc, m

    def _loadboat(self, who):
        """Load the boat. The source is implicit based on boat location."""

        if len(who) > self.boatcapacity:
            raise ValueError(f"boat capacity ({self.boatcapacity}) exceeded.")
        src = self.locations[self.boatlocation]
        for p in who:
            if p not in src:
                raise ValueError(f"{p} is not at {self.boatlocation}")

        self.locations[self.BOAT] = set(who)
        for p in who:
            src.remove(p)

        self.raise_violations(src, self.locations[self.BOAT])

    def _unloadboat(self, dest):
        self.boatlocation = dest
        dstkps = self.locations[dest]
        for x in self.locations[self.BOAT]:
            dstkps.add(x)

        self.raise_violations(dstkps)
        self.locations[self.BOAT] = set()

    def _nextloc(self, n):
        n += 1
        if n == len(self.locations):
            n = 1
        return n

    def transfer(self, who, dest=None):
        """Transfer everyone in who to dest.

        The source is implicit (the boat location).
        If no dest is given, the "next" destination is implied.
        """

        if dest is None:
            dest = self._nextloc(self.boatlocation)

        self.movetrail.append(self.move2str(who, dest))
        self._loadboat(who)
        self._unloadboat(dest)

    def move(self, s):
        """String 's' is like K1P1 and means move Knight 1 and Page 1.

        If there are multiple possible destinations, add /n e.g.,
               'K1P1/2' to move to destination #2 (dest ranges 1 .. n)

        Easier-for-humans version of transfer(). See that for details.
        """

        s = s.upper()
        who = []
        dest = None
        while s:
            match s[0], int(s[1]):
                case ('K', id):
                    who.append(self.knight_by_id(id))
                case ('P', id):
                    who.append(self.page_by_id(id))
                case ('/', dest):
                    pass
                case _:
                    raise ValueError(s)
            s = s[2:]
        self.transfer(who, dest)

    def move2str(self, who, dest=None):
        s = ""
        for x in who:
            s += "{}{}".format(x.typechar, x.id)
        if dest is not None:
            s += "/{}".format(dest)
        return s


def samplesolution():
    return checksolution(Puzzle(),
                         ['k1p1',
                          'k1',
                          'p0p2',
                          'p1',
                          'k0k2',
                          'k0p0',
                          'k0k1',
                          'p2',
                          'p0p1',
                          'p0',
                          'p0p2'])


def checksolution(z, sv):
    for m in sv[:-1]:
        z.move(m)
        if z.endstate:
            raise ValueError("premature end state reached")

    # the last move finishes the puzzle
    z.move(sv[-1])
    if not z.endstate:
        raise ValueError("didn't finish")
    return z


def islandsolution():
    return checksolution(Puzzle(kps=4, nlocations=3),
                         ['K3P3/2',
                          'K3/1',
                          'K2P2/3',
                          'P2/2',
                          'P3/1',
                          'P1P3/2',
                          'P1/1',
                          'K1P1/3',
                          'K2/1',
                          'K2K3/3',
                          'K2/1',
                          'K0P0/3',
                          'K3/1',
                          'K2K3/3',
                          'P0/2',
                          'P0P2/3',
                          'P0/2',
                          'P0P3/3'])


def islandsolution2():
    return checksolution(Puzzle(kps=4, nlocations=3),
                         ['P0P3/2',
                          'P0/1',
                          'P0P2/2',
                          'P0/1',
                          'P0P1/2',
                          'P0/1',
                          'P0K0/3',
                          'P0/2',
                          'P1/1',
                          'P1K1/3',
                          'P1/2',
                          'P2/1',
                          'P2K2/3',
                          'P2/2',
                          'P3/1',
                          'P3K3/3',
                          'P3/2',
                          'P3P2/3',
                          'P3/2',
                          'P3P1/3',
                          'P3/2',
                          'P3P0/3'])


def islandsolution3():
    return checksolution(Puzzle(kps=4, nlocations=3),
                         ['K1P1/2',
                          'K1/1',
                          'K0P0/3',
                          'K0/1',
                          'K1K0/3',
                          'K1/1',
                          'P2K2/3',
                          'P2/2',
                          'P1/1',
                          'K3K1/3',
                          'P0/1',
                          'P0P3/3',
                          'K1/1',
                          'K1P1/3',
                          'P0/2',
                          'P2P0/3'])


def samplesolution2():
    return checksolution(Puzzle(),
                         ['p0p1',
                          'p0',
                          'p0p2',
                          'p0',
                          'k1k2',
                          'p1k1',
                          'k0k1',
                          'p2',
                          'p0p1',
                          'p0',
                          'p0p2'])


def samplebad():
    # this will raise an exception for an unprotected Page
    z = Puzzle()
    z.move('k1')


def solver(z=None):
    if z is None:
        z = Puzzle()

    # Breadth-first search.
    #
    # Each new proposed state is FIFO queued so that it won't be examined
    # until all prior proposed states have been tried. Those states, in turn,
    # may of course also queue more future proposed states. In this way,
    # the solution with the fewest moves (least "depth") will be found
    # if any solutions exist at all.

    # statetrail: INFINITE LOOP DETECTION, PLUS QUEUE OPTIMIZATION
    #
    # the statetrail avoids re-examining states that are superficially
    # different but represent a logical state already examined.
    # For example:
    #      z1 = Puzzle()
    #      z2 = Puzzle()
    #      z1.move('P0')
    #      z2.move('P1')
    #
    # The moves for z1 and z2 are different, but the "logical state of the
    # puzzle" (one page moved to the next location) is identical for both.
    # The statetrail allows filtering these cases out. This is not just
    # an optimization but is also necessary for avoiding infinite loops:
    #      z1 = Puzzle()
    #      z1.move('P0')
    #      z1.move('P0')
    # In this (silly/trivial) sequence, the puzzle has ended up back in its
    # original state. The algorithm has to know that it has seen this state
    # already and not process it further. There are (obviously) much more
    # elaborate sequences of moves that will have this same looping property.
    # The statetrail filters those out as they occur.

    statetrail = {z.canonicalstate()}
    q = [z]               # queue of states to examine, starting with initial
    while q:
        z = q.pop(0)
        for loc, m in z.unique_legal_moves():
            z2 = z.clone()
            z2.transfer(m, loc)
            z2state = z2.canonicalstate()
            if z2state not in statetrail:
                if z2.endstate:
                    return z2
                statetrail.add(z2state)
                q.append(z2)

    # no solution was found
    return None