Solving logic puzzles with Z3

poker.py

#!/usr/bin/env python3
from z3 import (
    Solver,
    Function,
    EnumSort,
    IntSort,
    Const,
    Distinct,
    And,
    Or,
    sat,
    unsat,
)

ORDINALS = {
    1: "1st",
    2: "2nd",
    3: "3rd",
    4: "4th",
    5: "5th",
}


def main():
    solver = Solver()

    Name, name_consts = EnumSort("Name",
                                 ["Usain", "Terry", "Oliver", "Neil", "Ian"])
    usain, terry, oliver, neil, ian = name_consts

    Card, card_consts = EnumSort("Name",
                                 ["2", "5", "6", "7", "J", "Q", "K"])
    card_2, card_5, card_6, card_7, card_J, card_Q, card_K = card_consts

    left = Function("left", Name, Card)
    solver.add(Distinct([left(name) for name in name_consts]))
    for name in name_consts:
        solver.add(Or(left(name) == card_6,
                      left(name) == card_7,
                      left(name) == card_J,
                      left(name) == card_Q,
                      left(name) == card_K))

    right = Function("right", Name, Card)
    solver.add(Distinct([right(name) for name in name_consts]))
    for name in name_consts:
        solver.add(Or(right(name) == card_2,
                      right(name) == card_5,
                      right(name) == card_J,
                      right(name) == card_Q,
                      right(name) == card_K))

    fold = Function("fold", Name, IntSort())
    solver.add(Distinct([fold(name) for name in name_consts]))
    for name in name_consts:
        solver.add(fold(name) >= 1, fold(name) <= 5)

    # ... nobody was dealt a pair in hand.
    for name in name_consts:
        solver.add(left(name) != right(name))

    # One player, first to fold, ... "That's the worst possible hand..."
    name_clue2 = Const("name_clue2", Name)
    solver.add(fold(name_clue2) == 1)
    solver.add(Or(And(left(name_clue2) == card_2,
                      right(name_clue2) == card_7),
                  And(left(name_clue2) == card_7,
                      right(name_clue2) == card_2)))

    # The card in Oliver's right hand had given him a pair on the flop
    solver.add(right(oliver) == card_Q)

    # Ian folded 4th, having chased an inside straight that he'd seen on the
    # flop, though he didn't get so much as a pair
    solver.add(fold(ian) == 4)
    solver.add(left(ian) != card_J, right(ian) != card_J)
    solver.add(left(ian) != card_Q, right(ian) != card_Q)
    solver.add(left(ian) != card_K, right(ian) != card_K)

    # Terry had two pair with the turn card
    solver.add(Or(left(terry) == card_Q, right(terry) == card_Q))
    solver.add(Or(left(terry) == card_J, right(terry) == card_J))

    # The third player to fold, ... had a queen in his left hand
    name_clue6 = Const("name_clue6", Name)
    solver.add(fold(name_clue6) == 3)
    solver.add(left(name_clue6) == card_Q)

    # Neil folded next after Usain, refusing to chase an inside straight...
    solver.add(fold(neil) == fold(usain) + 1)

    if solver.check() == sat:
        m = solver.model()
        for name in name_consts:
            print("{} had {}{} and folded {}"
                  .format(name,
                          m.eval(left(name)),
                          m.eval(right(name)),
                          ORDINALS[m.eval(fold(name)).as_long()]))

        expressions = []
        for name in name_consts:
            expressions.append(left(name) != m.eval(left(name)))
            expressions.append(right(name) != m.eval(right(name)))
            expressions.append(fold(name) != m.eval(fold(name)))
        solver.add(Or(expressions))
        if solver.check() == unsat:
            print("Solution is unique")
        else:
            print("Solution is not unique")


if __name__ == "__main__":
    main()

requirements.txt

z3-solver

skiing.py

#!/usr/bin/env python3
from z3 import (
    Solver,
    Function,
    EnumSort,
    IntSort,
    RealSort,
    Const,
    Consts,
    Distinct,
    Or,
    Xor,
    sat,
    unsat,
)

"""
points: 82, 89, 96, 103
jumpers: Denise, Madeline, Patti, Shawna
distances: 90.1 meters, 95.0 meters, 96.3 meters, 102.9 meters

1. The skier who jumped 90.1 meters scored more points than Denise.
2. Patti scored 7 fewer points than the contestant who jumped 102.9 meters.
3. The jumper who scored 103 points was either Patti or the contestant who
jumped 102.9 meters.
4. Shawna jumped 96.3 meters.
5. Denise scored 89 points.
"""


def main():
    # create an EnumSort for names
    # store points/distances as ints/floats (so we can do relations later)
    # make functions from name to points and name to distance
    # handle clues that don't use names by making up a constant for it
    solver = Solver()

    Name, name_consts = EnumSort("Name",
                                 ["Denise", "Madeline", "Patti", "Shawna"])
    denise, madeline, patti, shawna = name_consts

    points = Function("points", Name, IntSort())
    solver.add(Distinct([points(name) for name in name_consts]))
    for name in name_consts:
        solver.add(Or(points(name) == 82,
                      points(name) == 89,
                      points(name) == 96,
                      points(name) == 103))

    distance = Function("distance", Name, RealSort())
    solver.add(Distinct([distance(name) for name in name_consts]))
    for name in name_consts:
        solver.add(Or(distance(name) == 90.1,
                      distance(name) == 95.0,
                      distance(name) == 96.3,
                      distance(name) == 102.9))

    # 1. The skier who jumped 90.1 meters scored more points than Denise.
    name_clue1 = Const("name_clue1", Name)
    solver.add(distance(name_clue1) == 90.1)
    solver.add(points(name_clue1) > points(denise))

    # 2. Patti scored 7 fewer points than the contestant who jumped 102.9
    # meters.
    name_clue2 = Const("name_clue2", Name)
    solver.add(points(patti) == points(name_clue2) - 7)
    solver.add(distance(name_clue2) == 102.9)

    # 3. The jumper who scored 103 points was either Patti or the contestant
    # who jumped 102.9 meters.
    name1_clue3, name2_clue3 = Consts("name1_clue3 name2_clue3", Name)
    solver.add(points(name1_clue3) == 103)
    solver.add(distance(name2_clue3) == 102.9)
    solver.add(Xor(name1_clue3 == patti, name1_clue3 == name2_clue3))

    # 4. Shawna jumped 96.3 meters.
    solver.add(distance(shawna) == 96.3)

    # 5. Denise scored 89 points.
    solver.add(points(denise) == 89)

    if solver.check() == sat:
        m = solver.model()
        for name in name_consts:
            print("{}: {} points, {}m"
                  .format(name,
                          m.eval(points(name)),
                          m.eval(distance(name)).as_decimal(1)))
        # eliminate this solution and check if it is unique
        expressions = []
        for name in name_consts:
            expressions.append(points(name) != m.eval(points(name)))
            expressions.append(distance(name) != m.eval(distance(name)))
        solver.add(Or(expressions))
        if solver.check() == unsat:
            print("Solution is unique")
        else:
            print("Solution is not unique")
    else:
        print("Contradiction!")


if __name__ == "__main__":
    main()

tv_ratings.py

#!/usr/bin/env python3
from z3 import (
    Solver,
    Function,
    EnumSort,
    IntSort,
    Const,
    Distinct,
    And,
    Or,
    Xor,
    sat,
    unsat,
)


def main():
    solver = Solver()

    Name, name_consts = EnumSort("Name",
                                 ["BNRG", "CVT", "KWTM", "PCR", "TWL"])
    bnrg, cvt, kwtm, pcr, twl = name_consts

    Show, show_consts = EnumSort("Show",
                                 ["Moneygab", "Ponyville", "Powertrips",
                                  "Soap_Suds", "Top_Chow"])
    moneygab, ponyville, powertrips, soapsuds, topchow = show_consts

    show = Function("show", Name, Show)
    solver.add(Distinct([show(name) for name in name_consts]))
    for name in name_consts:
        solver.add(Or(show(name) == moneygab,
                      show(name) == ponyville,
                      show(name) == powertrips,
                      show(name) == soapsuds,
                      show(name) == topchow))

    viewers = Function("viewers", Name, IntSort())
    solver.add(Distinct([viewers(name) for name in name_consts]))
    for name in name_consts:
        solver.add(viewers(name) >= 1, viewers(name) <= 5)

    channel = Function("channel", Name, IntSort())
    solver.add(Distinct([channel(name) for name in name_consts]))
    for name in name_consts:
        solver.add(Or(channel(name) == 15,
                      channel(name) == 22,
                      channel(name) == 43,
                      channel(name) == 56,
                      channel(name) == 62))

    # 1. Soap Suds (which doesn't have the most viewers) doesn't air on BNRG,
    # which has one million fewer viewers than the channel that airs
    # Powertrips.
    name1_clue1 = Const("name1_clue1", Name)
    name2_clue1 = Const("name2_clue1", Name)
    solver.add(show(name1_clue1) == soapsuds)
    solver.add(viewers(name1_clue1) != 5)
    solver.add(name1_clue1 != bnrg)
    solver.add(show(name2_clue1) == powertrips)
    solver.add(viewers(bnrg) == viewers(name2_clue1) - 1)

    # 2. Channel 22 has fewer viewers than the channel that airs Top Chow,
    # which has fewer viewers than TWL.
    name1_clue2 = Const("name1_clue2", Name)
    solver.add(channel(name1_clue2) == 22)
    name2_clue2 = Const("name2_clue2", Name)
    solver.add(show(name2_clue2) == topchow)
    solver.add(viewers(name1_clue2) < viewers(name2_clue2))
    solver.add(viewers(name2_clue2) < viewers(twl))

    # 3. TWL, which isn't carried on channel 56, boasts the most viewers of
    # all five channels.
    solver.add(channel(twl) != 56)
    solver.add(viewers(twl) == 5)

    # 4. Between PCR and the channel that airs Moneygab, one has three million
    # viewers and the other is on channel 22.
    name1_clue4 = Const("name1_clue4", Name)
    solver.add(show(name1_clue4) == moneygab)
    solver.add(Xor(And(viewers(pcr) == 3, channel(name1_clue4) == 22),
                   And(channel(pcr) == 22, viewers(name1_clue4) == 3)))

    # 5. Of channel 15 and the station that airs Ponyville, one is KWTM and
    # the other has the smallest viewership of all five channels.
    name1_clue5 = Const("name1_clue5", Name)
    solver.add(channel(name1_clue5) == 15)
    name2_clue5 = Const("name2_clue5", Name)
    solver.add(show(name2_clue5) == ponyville)
    solver.add(Xor(And(name1_clue5 == kwtm, viewers(name2_clue5) == 1),
                   And(name2_clue5 == kwtm, viewers(name1_clue5) == 1)))

    # 6. Moneygab doesn't air on either channel 15 or 56.
    name_clue6 = Const("name_clue6", Name)
    solver.add(show(name_clue6) == moneygab)
    solver.add(channel(name_clue6) != 15)
    solver.add(channel(name_clue6) != 56)

    # 7. CVT isn't carried on channel 62.
    solver.add(channel(cvt) != 62)

    if solver.check() == sat:
        m = solver.model()
        for name in name_consts:
            print("{}: {} million, {}, #{}"
                  .format(name,
                          m.eval(viewers(name)),
                          m.eval(show(name)),
                          m.eval(channel(name))))
        # eliminate this solution and check if it is unique
        expressions = []
        for name in name_consts:
            expressions.append(viewers(name) != m.eval(viewers(name)))
            expressions.append(channel(name) != m.eval(channel(name)))
            expressions.append(show(name) != m.eval(show(name)))
        solver.add(Or(expressions))
        if solver.check() == unsat:
            print("Solution is unique")
        else:
            print("Solution is not unique")
    else:
        print("Contradiction!")


if __name__ == "__main__":
    main()