Program to find number of sublists with sum k in a binary list in Python

Given a binary list containing only 0s and 1s, we need to find the number of contiguous sublists whose sum equals a given value k. This problem can be efficiently solved using the prefix sum technique with a hash map to track cumulative sums.

Problem Understanding

For a binary list nums = [1, 0, 0, 1, 1, 1, 0, 1] and k = 3, we need to find all contiguous sublists that sum to 3. The valid sublists are:

• [1,0,0,1,1] starting at index 0
• [0,0,1,1,1] starting at index 1
• [0,1,1,1] starting at index 2
• [1,1,1] starting at index 3
• [1,1,0,1] starting at index 4

And several other combinations, totaling 8 sublists.

Algorithm

The solution uses a prefix sum approach with the following steps:

• Maintain a dictionary sums to store frequency of prefix sums
• Initialize sums[0] = 1 to handle sublists starting from index 0
• Track running sum r_sum as we iterate through the array
• For each element, check if (r_sum - k) exists in our dictionary
• If it exists, add its frequency to our answer (these represent valid sublists)

Implementation

def solve(nums, k):
    sums = {0: 1}  # Initialize with 0 sum occurring once
    r_sum = 0      # Running sum
    ans = 0        # Count of valid sublists
    
    for x in nums:
        r_sum += x
        # Check if (r_sum - k) exists in sums dictionary
        ans += sums.get(r_sum - k, 0)
        # Update frequency of current running sum
        sums[r_sum] = sums.get(r_sum, 0) + 1
    
    return ans

# Test with the given example
nums = [1, 0, 0, 1, 1, 1, 0, 1]
k = 3
result = solve(nums, k)
print(f"Number of sublists with sum {k}: {result}")

Number of sublists with sum 3: 8

How It Works

The key insight is that if we have prefix sums prefix[i] and prefix[j] where prefix[j] - prefix[i] = k, then the sublist from index i+1 to j has sum k.

As we traverse the array:

• We calculate the running sum up to current position
• We check if (current_sum - k) has been seen before
• If yes, it means there are sublists ending at current position with sum k
• We add the frequency of (current_sum - k) to our answer

Example with Different Input

# Test with a smaller example
nums = [1, 1, 0, 1]
k = 2

result = solve(nums, k)
print(f"Binary list: {nums}")
print(f"Target sum: {k}")
print(f"Number of sublists with sum {k}: {result}")

# Let's trace through manually
print("\nSublists with sum 2:")
print("[1, 1] at indices 0-1")
print("[1, 0, 1] at indices 1-3")

Binary list: [1, 1, 0, 1]
Target sum: 2
Number of sublists with sum 2: 2

Sublists with sum 2:
[1, 1] at indices 0-1
[1, 0, 1] at indices 1-3

Time and Space Complexity

• Time Complexity: O(n) where n is the length of the input list
• Space Complexity: O(n) for storing prefix sum frequencies in the dictionary

Conclusion

This prefix sum approach efficiently finds all contiguous sublists with a given sum in O(n) time. The key insight is using a hash map to track previously seen prefix sums, allowing us to identify valid sublists in a single pass through the array.