Implementation for https://gist.github.com/hwayne/e5a65b48ab50a2285de47cfc11fc955f

budget.py

"""
You are writing simple budgeting softare. A Budget consists of a total limit,
represented as a positive integer, and a dictionary of category limits,
represented as a map from strings to integers. For example:

    {
        "total_limit": 50,
        "category_limits": {
            "food": 10,
            "rent": 11,
            "candles": 49
        }
    }

(You can assume that no category will have a higher limit than the total,
but you don't have to enforce this.)

A Bill is an array of Items. Each Item has a cost (positive integer), an optional
count (positive integer), and optional categories (array of strings). For example:

    [
        {"cost": 5},
        {"cost": 1, "count": 3},
        {"cost": 2, "categories": ["food", "gym"]},
        {"cost": 1, "count": 2, "categories": ["transit"]}
    ]

The goal is to write can_afford(budget, bill).

You do not need to use the same names for anything, or the same data layout,
as long as the requirements are satisfied.

Challenge from https://gist.github.com/hwayne/e5a65b48ab50a2285de47cfc11fc955f
"""

from copy import deepcopy
from itertools import chain, repeat
from typing import Dict, List
from typing_extensions import TypedDict


class Budget(TypedDict):
    total_limit: int
    category_limits: Dict[str, int]


class Item(TypedDict, total=False):
    cost: int  # Required
    count: int
    categories: List[str]


def can_afford(budget: Budget, bill: List[Item]) -> bool:
    """Affordability is calculated under the following rules:

    - The cost of every item is subtracted from total_limit.
    - If an item shares a category with the budget's category_limits, then price of
      that item is also subtracted from said category in addition to the total_limit.
    - The item might not share a category with the budget, or it might not have a
      category field at all. In both cases, no special rules apply.
    - If an item shares multiple categories with the budget, the cost is applied
      to only one category limit of your choice.
    - If an item has a count, the cost of the item is included count times.
      This applies to both total_limit and the category limits.
      If the count field is missing, you can assume a count of 1.

    can_afford is true if, after all items are accounted for, total_limit and all category
    limits are non-negative. If the total_limit or any category limits are below 0, then
    can_afford is false.

    can_afford should be as permissive as possible.
    """
    total_limit = budget["total_limit"]
    category_limits = dict(budget["category_limits"])
    multi_category = []

    # "count" is a pointless complication, so let's erase it from our datamodel:
    flat_bill = chain(*(repeat(x, x.pop("count", 1)) for x in deepcopy(bill)))

    # Implemented in two passes: first, we handle the simple cases with no discretion
    # about which category (if any) to charge; and then we go back for the interacting
    # items.  If at any point we're over budget, return False immediately.
    for item in flat_bill:
        total_limit -= item["cost"]
        if total_limit < 0:
            return False
        item_cats = set(item.get("categories", [])).intersection(category_limits)
        if len(item_cats) == 1:
            (cat,) = item_cats
            category_limits[cat] -= item["cost"]
            if category_limits[cat] < 0:
                return False
        elif len(item_cats) >= 2:
            item["categories"] = sorted(item_cats)
            multi_category.append(item)

    # If the simple first-pass handled every item, we're done.
    assert total_limit >= 0
    if not multi_category:
        return True

    # We now have a list of items which could be charged against any of several
    # categories, and have to find a maximally-permissive set of assignments.
    #
    # This sounds a lot like a Knapsack problem, so I'll just use a greedy
    # algorithm and let the property-testing framework show me counterexamples.
    for cost, cats in sorted(
        ((item["cost"], item["categories"]) for item in multi_category),
        reverse=True,
    ):
        highest_remaining_cat = max(cats, key=category_limits.__getitem__)
        category_limits[highest_remaining_cat] -= cost
        if category_limits[highest_remaining_cat] < 0:
            return False

    assert min(category_limits.values()) >= 0
    return True


# While the sole item in the bill is greater than category_limits[a],
# it is less than category_limits[b], so the bill is affordable.
assert can_afford(
    {"total_limit": 5, "category_limits": {"a": 1, "b": 3}},
    [{"cost": 2, "categories": ["a", "b"]}],
)

# Now, let's test with properties!
from hypothesis import event, given, settings, strategies as st


# We'll set some arbitrary limits - three known and one unknown category,
# prices in [1..10], budgets in [1..100], counts in [1..5] - just to ensure
# that we generate plenty of overlap while being expressive enough to find
# most bugs.  (If I'm wrong, that's an interesting data point!)
_categories = ("food", "rent", "misc", "unknown")
budgets: st.SearchStrategy[Dict[str, object]] = st.fixed_dictionaries(
    {
        "total_limit": st.integers(1, 100),
        "category_limits": st.dictionaries(
            keys=st.sampled_from(_categories[:-1]),
            values=st.integers(1, 100),
        ),
    }
)
bills: st.SearchStrategy[List[Dict[str, object]]] = st.lists(
    st.fixed_dictionaries(
        {"cost": st.integers(1, 10)},
        optional={
            "count": st.integers(1, 5),
            "categories": st.lists(st.sampled_from(_categories), unique=True),
        },
    ),
    min_size=1,
)


@settings(max_examples=10000)
@given(budgets, bills)
def test_can_afford(budget, bill):
    allowed = can_afford(budget, bill)
    event(f"allowed: {allowed}")

    total_should_be_allowed = budget["total_limit"] >= sum(
        item["cost"] * item.get("count", 1) for item in bill
    )
    if allowed:
        assert total_should_be_allowed
    elif total_should_be_allowed:
        # Instead of checking whether there's any satisfying assignment, we check
        # whether a few simple assignments of items to categories are satisfying.
        # Specifically, it must be impossible to assign every item to the n'th
        # category (of those in the budget) it could be placed in.
        assert budget["category_limits"]
        for n in (1, 2, 3, 4, 5):
            cats = dict(budget["category_limits"])
            for item in bill:
                allowed_cats = [c for c in item.get("categories", []) if c in cats]
                if allowed_cats:
                    c = allowed_cats[n % len(allowed_cats)]
                    cats[c] -= item["cost"] * item.get("count", 1)
            assert min(cats.values()) < 0


@given(budgets)
def test_can_always_afford_a_bill_with_no_items(budget):
    # This kind of thing is a classic trap, and it's easier to reason about
    # the performance of the main test if we handle empty bills separately.
    assert can_afford(budget, [])

As an author of Hypothesis I'd usually advocate for between 5% and 80% property-based tests depending on the project and domain, but in this case I'm happy with 2 PBTs and zero example-based tests (aside from the spec example, if that counts).

I don't expect that these two tests can catch bugs in all implementations, but they're enough of a demo for this exercise.

For a serious project I'd want to actually prove that the greedy approach works; but for something quick-and-dirty generating a few million examples where it doesn't fail on a focussed test is a good start! It's at least not obviously-broken before starting a proof 🙂

Note that it's trivial for the test I actually wrote to find a case where the greedy approach fails, just by tuning the strategies a bit: require that budgets have a limit for each category and that all items in a bill have two or more categories, and you promptly get this counterexample:

Falsifying example: test_can_afford(
    bill=[
        {"categories": ["food", "rent", "misc"], "cost": 3},
        {"categories": ["food", "rent"], "cost": 2},
    ],
    budget={"category_limits": {"food": 1, "misc": 3, "rent": 4}, "total_limit": 5},
)

It's affordable if we charge the first item to misc and the second to rent; but the greedy approach charges the first to rent (higher remaining budget) and then can't afford the second.