diff --git a/knapsack/knapsack.py b/knapsack/knapsack.py index 0648773c919f..fb89c1a4f194 100644 --- a/knapsack/knapsack.py +++ b/knapsack/knapsack.py @@ -4,7 +4,7 @@ from __future__ import annotations -from functools import lru_cache +from functools import cache def knapsack( @@ -12,7 +12,7 @@ def knapsack( weights: list[int], values: list[int], counter: int, - allow_repetition=False, + allow_repetition: bool = False, ) -> int: """ Returns the maximum value that can be put in a knapsack of a capacity cap, @@ -37,7 +37,7 @@ def knapsack( got the weight of 10*5 which is the limit of the capacity. """ - @lru_cache + @cache def knapsack_recur(capacity: int, counter: int) -> int: # Base Case if counter == 0 or capacity == 0: @@ -62,6 +62,65 @@ def knapsack_recur(capacity: int, counter: int) -> int: return knapsack_recur(capacity, counter) +def knapsack_with_count( + capacity: int, + weights: list[int], + values: list[int], + counter: int, + allow_repetition: bool = False, +) -> tuple[int, int]: + """ + Return both the maximum knapsack value and number of optimal selections. + + The return value is ``(max_value, number_of_optimal_selections)``. + If multiple choices produce the same maximum value, their counts are added. + Distinct selections are order-insensitive: + - with ``allow_repetition=False`` these are distinct subsets by item index; + - with ``allow_repetition=True`` these are distinct multisets by item index + multiplicity. + + >>> cap = 50 + >>> val = [60, 100, 120] + >>> w = [10, 20, 30] + >>> c = len(val) + >>> knapsack_with_count(cap, w, val, c) + (220, 1) + >>> knapsack_with_count(cap, w, val, c, True) + (300, 1) + >>> knapsack_with_count(3, [1, 2, 3], [1, 2, 3], 3) + (3, 2) + >>> knapsack_with_count(2, [1, 2], [1, 2], 2, True) + (2, 2) + """ + + @cache + def knapsack_recur(remaining_capacity: int, item_count: int) -> tuple[int, int]: + # Base Case: one empty subset yields value 0. + if item_count == 0 or remaining_capacity == 0: + return 0, 1 + + if weights[item_count - 1] > remaining_capacity: + return knapsack_recur(remaining_capacity, item_count - 1) + + left_capacity = remaining_capacity - weights[item_count - 1] + included_value, included_count = knapsack_recur( + left_capacity, item_count if allow_repetition else item_count - 1 + ) + included_value += values[item_count - 1] + + excluded_value, excluded_count = knapsack_recur( + remaining_capacity, item_count - 1 + ) + + if included_value > excluded_value: + return included_value, included_count + if excluded_value > included_value: + return excluded_value, excluded_count + return included_value, included_count + excluded_count + + return knapsack_recur(capacity, counter) + + if __name__ == "__main__": import doctest diff --git a/knapsack/tests/test_knapsack.py b/knapsack/tests/test_knapsack.py index 80378aae4579..5f7ba231b59f 100644 --- a/knapsack/tests/test_knapsack.py +++ b/knapsack/tests/test_knapsack.py @@ -58,6 +58,26 @@ def test_knapsack_repetition(self): c = len(val) assert k.knapsack(cap, w, val, c, True) == 300 + def test_knapsack_with_count(self): + """ + test for maximum value and number of optimal selections + """ + cap = 50 + val = [60, 100, 120] + w = [10, 20, 30] + c = len(val) + assert k.knapsack_with_count(cap, w, val, c) == (220, 1) + assert k.knapsack_with_count(cap, w, val, c, True) == (300, 1) + assert k.knapsack_with_count(0, w, val, c) == (0, 1) + assert k.knapsack_with_count(50, w, val, 0) == (0, 1) + + def test_knapsack_with_count_ties(self): + """ + test tie handling for counting optimal subsets + """ + assert k.knapsack_with_count(3, [1, 2, 3], [1, 2, 3], 3) == (3, 2) + assert k.knapsack_with_count(2, [1, 2], [1, 2], 2, True) == (2, 2) + if __name__ == "__main__": unittest.main()