Subset Sum using mknapsack library

subset_sum.py

import numpy as np
from mknapsack import solve_subset_sum


def subset_sum(*, weights: np.array, capacity: float) -> np.array:
    assert capacity > 0, f"Capacity must be positive!, given: {capacity}"

    n_weights = len(weights)

    if n_weights < 2:
        return np.zeros(shape=0, dtype=np.uint8) if n_weights == 0 else np.ones(shape=1, dtype=np.uint8)

    weights_to_consider_bool = (weights > 0) & (weights <= capacity)
    weights_to_consider = weights[weights_to_consider_bool]
    weights_sum = weights_to_consider.sum()

    if weights_sum == capacity:
        return weights_to_consider_bool.astype(np.uint8)

    if (len(weights_to_consider) < 2) or (weights_sum < capacity):
        abs_diff = np.abs(capacity - weights)
        diff1 = abs_diff.min()
        diff2 = capacity - weights_sum

        if diff1 <= diff2:
            ret = np.zeros(shape=n_weights, dtype=np.uint8)
            ret[np.where(abs_diff == diff1)] = 1
            return ret

        return weights_to_consider_bool.astype(np.uint8)

    ret = np.zeros(shape=n_weights, dtype=np.uint8)
    ret[weights_to_consider_bool] = solve_subset_sum(
        weights=weights_to_consider,
        capacity=capacity,
        method_kwargs={"check_inputs": False},
    )
    return ret


## Usage-
## weights = np.array(choices)
## capacity = target
## bool_arr = subset_sum(weights=weights, capacity=capacity)
## # bool_arr will be a binary(0/1) vector 1 denotating the weight[i] to condier for packing
## bool_arr.dot(weights) ~ capacity
## This code will try to find the closest combination of weights to satisfy the given target
## Note that the total sum of weights considered by bool_arr might be slightly greater than the given capacity