Knapsack 0-1 solver

bin_knapsack.py

"""Binary Knapsack solver"""

from sys import exit as sys_exit


def ask_number(prompt: str) -> int | None:
    """Asks an integer number repeatedly"""
    while True:
        try:
            return int(input(prompt).strip())
        except ValueError:
            continue
        except KeyboardInterrupt:
            return None


def ask_num_vector(n: int) -> list[int] | None:
    """Asks integer values for each i-th object, having `n` total objects"""
    vec = []
    for i in range(n):
        p = ask_number(f"\nInsert value for the object x{i+1}\n>>")
        if p is None:
            return None
        if p < 0:
            print("\nProfit value must be >= 0")
            return None
        vec.append(p)
    return vec


def print_matrix(matrix: list[list[int]]) -> None:
    """Pretty print a matrix"""
    for row in matrix:
        print(" ".join(str(element) for element in row))


def solve(
    profits_vec: list[int], weights_vec: list[int], n_objects: int, total_cap: int
) -> tuple[int, list[bool]] | None:
    """Solves knapsack 0-1
    Returns a two elements tuple, with the first element being the optimal value,
    the second element is the optimal (binary) solution.
    Returns None in case of errors...
    """

    matrix = [[None for _ in range(n_objects + 1)] for _ in range(total_cap + 1)]
    # Sets first row and column with 0
    for i in range(n_objects + 1):
        matrix[0][i] = 0
    for i in range(total_cap + 1):
        matrix[i][0] = 0

    # Finds the optimal value
    for k in range(1, total_cap + 1):
        for i in range(1, n_objects + 1):

            p = profits_vec[i - 1]
            w = weights_vec[i - 1]
            if k >= w:
                matrix[k][i] = max(matrix[k][i - 1], p + matrix[k - w][i - 1])
            else:
                matrix[k][i] = matrix[k][i - 1]

    print_matrix(matrix)

    # Optimal value found
    zstar = matrix[total_cap][n_objects]

    xstar = [False] * n_objects

    # Find optimal solution
    xstar[n_objects - 1] = bool(zstar)
    for i in range(n_objects - 1, 0, -1):
        if (
            total_cap >= weights_vec[i - 1]
            and matrix[total_cap][i] != matrix[total_cap - weights_vec[i - 1]][i - 1]
        ):
            total_cap -= weights_vec[i - 1]
            xstar[i - 1] = True

    return (zstar, xstar)


def __main(
    profits_vec: list[int] = None,
    weights_vec: list[int] = None,
    n_objects: int = None,
    total_cap: int = None,
):
    """Main entry point"""

    if n_objects is None:
        n_objects = ask_number("\nHow many objects?\n>>")
        if n_objects is None or n_objects <= 0:
            print("\nMust be > 0")
            return -1

    if total_cap is None:
        total_cap = ask_number("\nTotal capacity of the sack?\n>>")
        if total_cap is None or total_cap < 0:
            print("\nMust be >= 0")
            return -1

    if profits_vec is None:
        print("\nPlease insert the profit values for each i-th object...")
        profits_vec = ask_num_vector(n_objects)

    if weights_vec is None:
        print("\nPlease insert the weights values for each i-th object...")
        weights_vec = ask_num_vector(n_objects)

    zstar, xstar = solve(profits_vec, weights_vec, n_objects, total_cap)

    print(
        "\nz = " + " + ".join((f"{p}(x{i+1})" for i, p in enumerate(profits_vec))),
        end="\n\n",
    )
    print(
        " + ".join((f"{w}(x{i+1})" for i, w in enumerate(weights_vec)))
        + f" <= {total_cap}"
    )
    print("\nxi ∈ {0,1} i = {" + " ".join(str(i + 1) for i in range(n_objects)) + "}")

    print("\nSolution is  x* = |" + " ".join("1" if b else "0" for b in xstar) + "|")
    print(f"\nOptimal is  z* = {zstar}")

    return 0


if __name__ == "__main__":
    sys_exit(__main())