Faffing around with the game 'set'.

set.py

#!/usr/bin/env python3

from typing import Generator, Union, Iterable
import random
from itertools import repeat, product


class Set(frozenset):
    """
    A 'Set' of cards from the game set. This class is to avoid confusion with the
    standard python 'set' datastructure, which is also used here.
    """

    @classmethod
    def from_featurewise(cls, featurewise) -> "Set":
        """
        Take a valid set grouped by feature, e.g.:

            ((0, 0, 0), (1, 1, 1), (2, 2, 2), (0, 1, 2))

        And return the unique cards this set implies, i.e.:

            ((0, 1, 2, 0), (0, 1, 2, 1), (0, 1, 2, 2))
        """
        s = cls(zip(*featurewise))
        if len(s) == 3:  # sets must be len 3
            return s
        else:
            raise ValueError("invalid set")


# sets have (0, 0, 0), (1, 1, 1), (2, 2, 2) or (0, 1, 2)
# features in a set must all be the same, or all different
VALID_FEATURE_COMBINATIONS = (
    (0, 0, 0),
    (1, 1, 1),
    (2, 2, 2),
    (0, 1, 2),
    (0, 2, 1),
    (1, 0, 2),
    (1, 2, 0),
    (2, 1, 0),
    (2, 0, 1),
)


def generate_valid_sets() -> Generator[Set, None, None]:
    """
    featurewise contains tuples like:

      ((0, 0, 0), (0, 1, 2), (1, 1, 1), (2, 2, 2))

    Where the first group describe the first feature (here, all 0s), the second group
    describes the second feature (all different in ascending order). ('featurewise',
    because this consists of tuples of features.)
    """
    for featurewise in product(*repeat(VALID_FEATURE_COMBINATIONS, 4)):
        try:
            yield Set.from_featurewise(featurewise)
        except ValueError:
            continue


def sample_contains_set(sample) -> bool:
    """Test if a sample of cards contains a set"""
    return any(all(card in sample for card in SET) for SET in ALL_SETS)


def find_set(sample: Iterable[tuple[int, int, int, int]]) -> Union[Set, None]:
    """Find a set in a sample of cards"""
    for SET in ALL_SETS:
        if all(card in sample for card in SET):
            return SET


def random_sample_contains_set(k: int = 12) -> bool:
    """Test if a random sample of k cards contains a set"""
    return sample_contains_set(random.sample(tuple(ALL_CARDS), k=k))


def proportion_random_samples_contain_sets(n: int, k: 12) -> int:
    """Count of how many random samples of k cards contain sets"""
    return sum(random_sample_contains_set(k=k) for _ in range(n)) / n


ALL_SETS = frozenset(generate_valid_sets())
ALL_CARDS = frozenset(product(*repeat(range(3), 4)))


def play_game():
    deck = list(random.shuffle(ALL_CARDS))

    table = set()
    for _ in range(12):
        table.add(deck.pop())

    # find_set(table)


if __name__ == "__main__":
    random.seed(42)

    # Chance of there not being a set in 12 random cards is apparently ~33:1
    print(proportion_random_samples_contain_sets(n=10_000, k=12))
    print(proportion_random_samples_contain_sets(n=10_000, k=15))

test_set.py

from set import *

assert len(ALL_CARDS) == 81
for CARD in ALL_CARDS:
    assert isinstance(CARD, tuple)
    assert len(CARD) == 4
    assert set(CARD) - {0, 1, 2} == set()

assert Set(((0, 1, 2, 0), (1, 2, 1, 1), (2, 0, 0, 2))) in ALL_SETS

assert sample_contains_set(sample=((0, 0, 0, 0), (1, 1, 1, 1), (2, 2, 2, 2)))
assert sample_contains_set(
    sample=((0, 0, 0, 0), (1, 1, 1, 1), (2, 2, 2, 2), (0, 1, 2, 1))
)
assert sample_contains_set(sample=((0, 1, 2, 0), (1, 2, 1, 1), (2, 0, 0, 2)))
assert not sample_contains_set(sample=((1, 1, 1, 1), (2, 2, 2, 2), (0, 1, 2, 1)))