Find the orders with the best price for ordering a quantity from a set of offers

price_solver.py

# -*- coding: utf-8 -*-

from collections import defaultdict
import collections.abc
import itertools

from attr import field, frozen, mutable
import fastcore.basics
import more_itertools
import pydash

_ZERO_RET_FUNC = pydash.constant(0)


def create_orders(offers, request):
    offers = tuple(_filter_offers(offers.items()))
    return _backtrack(offers, request, _build_matrix(offers, request))


def print_orders(orders):
    total = _OrdersTotal()

    for cur_order in orders:
        total.update(cur_order)

    print(total)


@frozen
class QuantityOffer:
    price: float

    shop: object


@frozen
class OfferInfo:
    offer: QuantityOffer

    quantity: int


@frozen
class _BuyNode:
    offer_index: int

    price: float


@frozen
class _BuyMatrix:
    def __init__(self, offers, request):
        # Don't create a requests list for the last offer index if it'd
        # only contain one request.
        num_of_offers = min(len(offers), request - 1)
        # We use a placeholder buy node for all slots. It doesn't matter
        # what it contains since it's gonna be overwritten during the
        # initialization of the matrix. I'm using an invalid index so
        # that an exception may be raised if an attempt is made to
        # access this index.
        dummy_buy_node = _BuyNode(num_of_offers, 0)
        matrix = map(
            lambda offer_idx: [dummy_buy_node] * (request - offer_idx),
            range(1, num_of_offers),
        )
        # The first offer will have an extra slot for the original
        # request.
        offer0_reqs = request + 1
        self.__attrs_init__(
            tuple(
                more_itertools.prepend([dummy_buy_node] * offer0_reqs, matrix)
            ),
            offers,
        )

    def get(self, request_node):
        return (
            # The requested cell must have been previously filled.
            self._matrix[request_node.offer_index][request_node.request]
            if request_node.request > 0
            # Any index will work here as it isn't gonna be checked
            # anyway; only the price(for convenience) has to be zero.
            # I'm using an invalid index so that an exception may be
            # raised if an attempt is made to access this index.
            else _BuyNode(self.num_of_offers, 0)
        )

    def fill(self, request_node):
        num_of_offers = len(self._offers)
        self._matrix[request_node.offer_index][request_node.request] = min(
            map(
                lambda cur_offer: self._make_buy_node(
                    cur_offer,
                    request_node.request - self._offers[cur_offer].quantity,
                    self._offers[cur_offer].offer.price,
                ),
                range(request_node.offer_index, num_of_offers),
            ),
            key=fastcore.basics.Self.price(),
        )

    def _make_buy_node(self, offer_index, prev_request, offer_price):
        return _BuyNode(
            offer_index,
            offer_price
            + self.get(_RequestNode(offer_index, prev_request)).price,
        )

    @property
    def num_of_offers(self):
        return len(self._matrix)

    _matrix: tuple[list[_BuyNode], ...]

    _offers: collections.abc.Sequence[OfferInfo]


@mutable
class _MutableFloat:
    value: float


@mutable
class _OrdersTotal:
    def update(self, order):
        self._order_groups[order.offer.shop][order.quantity] += 1
        self._quantity += order.quantity
        self._price += order.offer.price

    _order_groups = field(
        factory=lambda: defaultdict(lambda: defaultdict(_ZERO_RET_FUNC)),
        init=False,
    )

    _price = field(default=0, init=False)

    _quantity = field(default=0, init=False)


@mutable
class _RequestNode:
    offer_index: int

    request: int


def _backtrack(offers, request, matrix):
    orders = []
    req_node = _RequestNode(0, request)

    while req_node.request > 0:
        _update_orders(
            offers, req_node, matrix.get(req_node).offer_index, orders
        )

    return orders


def _build_matrix(offers, request):
    matrix = _BuyMatrix(offers, request)
    num_of_offers = len(offers)
    # The matrix is divided into two parts. The first one is a
    # rectangular area where all offer indices are filled for the
    # requests in this part. The second is a triangular area where the
    # number of offer indices needing to be filled decreases gradually
    # as the request gets higher.
    # We skip initializing cells with request = 0 since it's
    # impractical. The buy matrix structure ensures that such cells are
    # never read.
    coord_range = itertools.product(
        range(1, request - num_of_offers + 1), range(matrix.num_of_offers)
    )
    _fill_rect(matrix, map(reversed, coord_range))
    coord_range = range(min(num_of_offers, request) - 1, 0, -1)
    _fill_triangle(
        matrix,
        itertools.starmap(
            lambda req_coord, offer_coord: (
                offer_coord - 1,
                request - req_coord,
            ),
            itertools.combinations_with_replacement(coord_range, 2),
        ),
    )
    matrix.fill(_RequestNode(0, request))  # the original request
    return matrix


def _fill_rect(matrix, cell_coords):
    # A full request is one that needs to be filled for all used offers.
    for cur_coord in cell_coords:
        matrix.fill(_RequestNode(*cur_coord))


def _fill_triangle(matrix, cell_coords):
    # An incomplete request is one that only needs to be filled for some
    # offers.
    for cur_coord in cell_coords:
        matrix.fill(_RequestNode(*cur_coord))


def _filter_offers(offers):
    # We sort the list of all offers by price. For ties with the same
    # price, we break by the quantity(descending). Specifying a negative
    # quantity is just one way to request a descending sort.
    return _filter_useless(
        sorted(
            _flatten_offers(offers),
            key=lambda offer_info: (
                offer_info.offer.price,
                -offer_info.quantity,
            ),
        )
    )


def _filter_useless(offers):
    # A useless offer is one that sells the same quantity(or lower) for
    # the same price(or higher). Given that the offers list is already
    # sorted by price(ascending) then quantity(descending), the sequence
    # resulting from dropping useless offers is already sorted by
    # quantity in ascending order also.
    qty = _MutableFloat(0)

    for cur_offer in offers:
        if cur_offer.quantity > qty.value:
            yield _update_qty(cur_offer, qty)


def _flatten_offers(offers):
    flattened = []

    for shop, catalog in offers:
        _update_offers(shop, catalog.items(), flattened)

    return flattened


def _update_offers(shop, catalog, offers):
    for qty, price in catalog:
        offers.append(OfferInfo(QuantityOffer(price, shop), qty))


def _update_orders(offers, request_node, offer_index, orders):
    orders.append(offers[offer_index])
    request_node.offer_index = offer_index
    request_node.request -= offers[offer_index].quantity


def _update_qty(offer, quantity):
    quantity.value = offer.quantity
    return offer