"Constrained Subsequence Sum" solution

src/constrained_subsequence_sum.py

@@ -0,0 +1,20 @@
+import heapq
+
+
+class Solution:
+    def constrainedSubsetSum(self, nums: list[int], k: int) -> int:
+        assert nums
+
+        result = nums[0]
+        heap = [(-nums[0], 0)]
+
+        for i in range(1, len(nums)):
+            while heap and i - heap[0][1] > k:
+                heapq.heappop(heap)
+
+            curr_max = max(nums[i], -heap[0][0] + nums[i])
+            result = max(result, curr_max)
+
+            heapq.heappush(heap, (-curr_max, i))
+
+        return result

tests/test_constrained_subsequence_sum.py

@@ -0,0 +1,15 @@
+import pytest
+
+from src.constrained_subsequence_sum import Solution
+
+
+@pytest.mark.parametrize(
+    "nums,k,expected",
+    (
+        ([10, 2, -10, 5, 20], 2, 37),
+        ([-1, -2, -3], 1, -1),
+        ([10, -2, -10, -5, 20], 2, 23),
+    ),
+)
+def test_solution(nums, k, expected):
+    assert Solution().constrainedSubsetSum(nums, k) == expected