Stolen Rubies riddle

stolen_rubies.py

#!/usr/bin/env python3

# Solves the stolen rubies riddle: https://www.youtube.com/watch?v=2QJ2L2ip32w

import itertools


def main():
    configs = unique_configs(gen_chests())
    # For each "reasonable" strategy (possibly assigning a different number to
    # each chest, where we determine the range based on the range of possible
    # chest denominations), we find the minimum score across all possible
    # configurations. This is the guaranteed output of that strategy. Note that
    # the intuitive thing to do is to assign each chest the _same_ number, since
    # we don't have any information to break the initial symmetry. However, it's
    # not immediately obvious that this is necessarily the case. To be safe, we
    # consider all possible strategies, which accounts for scenarios where you
    # hedge against some chests being more valuable than others.
    min_rubies = min(itertools.chain(*configs))
    max_rubies = max(itertools.chain(*configs))
    best_min_score = None
    best_strategy = None
    for strategy in gen_strategies(min_rubies, max_rubies):
        min_score = None
        for config in configs:
            for chests in itertools.permutations(config):
                # NOTE: This will produce duplicates if there are duplicate ruby
                # counts, but we tolerate that.
                score = compute_score(strategy, chests)
                if min_score is None or score < min_score:
                    min_score = score
        if best_min_score is None or min_score > best_min_score:
            best_min_score = min_score
            best_strategy = strategy
    print("best guanteed score:", best_min_score)
    print("best strategy:", best_strategy)


def compute_score(strategy, chests):
    score = 0
    for guess, rubies in zip(strategy, chests):
        if guess <= rubies:
            score += guess
    return score


def gen_strategies(min_rubies, max_rubies):
    for a in range(min_rubies, max_rubies + 1):
        for b in range(min_rubies, max_rubies + 1):
            for c in range(min_rubies, max_rubies + 1):
                yield (a, b, c)


def unique_configs(chest_tuples):
    # We create a unique representation for each multiset of ruby counts by
    # sorting by rubies. This allows us to avoid checking "identical" configs
    # multiple times. Since we are maximizing the minimum score over all
    # configurations, this is fine. If we needed to maximize expectation, we
    # would need a model of how likely each configuration was, and that's
    # unclear since it depends on the thief's behavior.
    def unique_rep(chests):
        return tuple(sorted(chests))

    return sorted(set(map(unique_rep, chest_tuples)))


def gen_chests():
    """Generate all possible arrangements of rubies in chests."""
    for n in range(2, 21):
        # The third box will have somewhere between 2 and 20 rubies (inclusive).
        # x is 6 less than y
        remaining = (30 - n) // 2
        x = remaining - 3
        y = remaining + 3
        # Deal with truncated division by ensuring this adds up.
        if n + x + y == 30:
            yield (n, x, y)


if __name__ == "__main__":
    main()