warehouse-solver.py

#!/usr/bin/env python3

def solve(customer, warehouses):
    """
    Choose warehoues to take products for customer base on some rules.

    Input:
        `customer`: Dictionary contains city & product list of a customer
        Eg: {"city": "Hanoi", "product": {"book": 1, "pen": 2}}

        `warehouses`: List of warehouses which are available to us
        Eg:
            [
                {"city": "Hanoi", "name": "A", "product": {"book": 1, "pen": 2}},
                {"city": "Hochiminh", "name": "B", "product": {"book": 10}},
            ]

    Output:
        Dictionary of warehouses with corresponding products
        Eg:
            {"A": {"book": 1, "pen": 2}}
    """
    customer_products = set(customer["product"].keys())

    remain_product_dict = customer["product"].copy()
    remain_product_count = sum(remain_product_dict.values())

    final_result = {}

    def sort_func(warehouse):
        """
        Function for sorting warehouse by priorities
        """
        rank = []

        # Priority #1: Same city
        if warehouse["city"] == customer["city"]:
            rank.append(1)
        else:
            rank.append(0)

        # Priority #2: Has all product
        if set(warehouse["product"].keys()) >= customer_products:
            rank.append(1)
        else:
            rank.append(0)

        # Priority #3: Largest quantity of a product, we only check products which
        # are in customer order, and only get the largest
        rank.append(max([v for k, v in warehouse["product"].items()
                         if k in customer_products]))

        return rank

    for warehouse in sorted(warehouses, key=sort_func, reverse=True):
        take_product = {}
        for product_name, product_count in warehouse["product"].items():
            if remain_product_count == 0:
                break

            if product_count == 0:
                continue

            if remain_product_dict[product_name] > 0:
                take_product[product_name] = min(remain_product_dict[product_name], product_count)

                remain_product_dict[product_name] -= take_product[product_name]
                warehouse["product"][product_name] -= take_product[product_name]
                remain_product_count -= take_product[product_name]

        if take_product:
            final_result[warehouse["name"]] = take_product

    return final_result


if __name__ == '__main__':
    customer = {"city": "Hanoi", "product": {"book": 2, "pen": 3}}

    A = {"city": "Hanoi", "name": "A", "product": {"book": 1}}
    B = {"city": "Hanoi", "name": "B", "product": {"book": 3, "pen": 3}}
    C = {"city": "Hanoi", "name": "C", "product": {"book": 10, "pen": 7}}
    D = {"city": "Hochiminh", "name": "D", "product": {"book": 30, "pen": 70}}
    assert str(solve(customer, [A, B, C, D])) == str({"C": {"book": 2, "pen": 3}})

    A = {"city": "Hanoi", "name": "A", "product": {"book": 1}}
    B = {"city": "Hochiminh", "name": "B", "product": {"book": 10}}
    C = {"city": "Hochiminh", "name": "C", "product": {"pen": 5}}
    assert str(solve(customer, [A, B, C])) == str({"A": {"book": 1}, "B": {"book": 1}, "C": {"pen": 3}})

    A = {"city": "Hanoi", "name": "A", "product": {"book": 1}}
    B = {"city": "Hanoi", "name": "B", "product": {"book": 3, "pen": 3, "paper": 100}}
    C = {"city": "Hanoi", "name": "C", "product": {"book": 10, "pen": 7}}
    D = {"city": "Hochiminh", "name": "D", "product": {"book": 30, "pen": 70}}
    assert str(solve(customer, [A, B, C, D])) == str({"C": {"book": 2, "pen": 3}})