Add knapsack variant to count optimal-value subsets

knapsack/knapsack.py

@@ -4,15 +4,15 @@
 
 from __future__ import annotations
 
-from functools import lru_cache
+from functools import cache
 
 
 def knapsack(
     capacity: int,
     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
 

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()