main.py

import numpy as np


def split_price(split, warehouse_prices, warehouse_min_price, freight_cost):
    price_per_warehouse = list(map(lambda row: sum(row[0] * row[1]), zip(warehouse_prices, split)))
    price = sum(price_per_warehouse)
    for warehouse_id, warehouse_price in enumerate(price_per_warehouse):
        if 0 < warehouse_price < warehouse_min_price[warehouse_id]:
            price += freight_cost
    return price


def best_split(blanks, warehouse_count, warehouse_prices, warehouse_quantities, warehouse_min_price, freight_cost):
    # order in which it would be best to take skus from warehouses
    sorted_warehouses = list(map(lambda blank_id: list(warehouse_prices[:, blank_id].argsort()), blanks))

    # calculate best possible split regardless of shipping
    split = np.zeros((warehouse_count, len(blanks)), dtype=np.int32)
    blanks_comp = blanks.copy()

    for blank_id in blanks:
        for warehouse_id in sorted_warehouses[blank_id]:
            wbc = min(warehouse_quantities[warehouse_id, blank_id], blanks_comp[blank_id])
            split[warehouse_id, blank_id] = wbc
            blanks_comp[blank_id] -= wbc
            if blanks_comp[blank_id] == 0:
                break

    price = split_price(split, warehouse_prices, warehouse_min_price, freight_cost)
    nonempty_warehouses = list(np.nonzero(split.sum(1))[0])

    current_price = price
    current_split = split
    # iterate - try to move items from one warehouse to another
    # break if there is no change in the last iteration
    while True:
        updated_price = False

        # Swap whole warehouse order
        for i, warehouse_id in enumerate(nonempty_warehouses):
            for next_warehouse_id in nonempty_warehouses[:i] + nonempty_warehouses[i + 1:]:
                row = np.copy(split[warehouse_id, :])
                split[warehouse_id, :] -= row
                split[next_warehouse_id, :] += row
                if (warehouse_quantities - split < 0).any():
                    split[next_warehouse_id, :] -= row
                    split[warehouse_id, :] += row
                    continue

                next_price = split_price(split, warehouse_prices, warehouse_min_price, freight_cost)
                if next_price < current_price:
                    current_price = next_price
                    current_split = np.copy(split)
                    updated_price = True

                split[next_warehouse_id, :] -= row
                split[warehouse_id, :] += row

        # Swap SKUs from one warehouse to another
        for i, warehouse_id in enumerate(nonempty_warehouses):
            for blank_id in range(len(blanks)):
                for next_warehouse_id in nonempty_warehouses[:i] + nonempty_warehouses[i + 1:]:
                    value = split[warehouse_id, blank_id]
                    split[warehouse_id, blank_id] -= value
                    split[next_warehouse_id, blank_id] += value
                    if (warehouse_quantities - split < 0).any():
                        split[next_warehouse_id, blank_id] -= value
                        split[warehouse_id, blank_id] += value
                        continue

                    next_price = split_price(split, warehouse_prices, warehouse_min_price, freight_cost)
                    if next_price < current_price:
                        current_price = next_price
                        current_split = np.copy(split)
                        updated_price = True

                    split[next_warehouse_id, blank_id] -= value
                    split[warehouse_id, blank_id] += value

        split = current_split
        nonempty_warehouses = list(np.nonzero(split.sum(1))[0])

        if not updated_price:
            break

    return current_price, split


def main():
    # sku_to_id = {'#1': 0, '#2': 1, '#3': 2}
    blanks = {0: 10, 1: 15, 2: 20}

    w_count = 4

    w_prices = np.array([[2, 3, 1], [2.5, 4, 2], [5, 1.5, 1.4], [1.9, 3.5, 1.23]])
    w_quantities = np.array([[1, 10, 15], [10, 4, 20], [20, 20, 7], [2, 7, 3]], dtype=np.int32)

    w_min_prices = [1, 300, 200, 30]
    freight_cost = 5

    price, split = best_split(blanks, w_count, w_prices, w_quantities, w_min_prices, freight_cost)
    print(price)
    print(split)


if __name__ == '__main__':
    main()