Prefix Sum

Prefix sum refers to the cumulative sum of elements in an array up to a specific index. It helps optimize computations involving cumulative sums. Let's see an example:


# Function to calculate prefix sum
def calculate_prefix_sum(array):
    prefix_sum = [0] * len(array)
    prefix_sum[0] = array[0]

    # Compute prefix sum for the remaining elements
    for i in range(1, len(array)):
        prefix_sum[i] = prefix_sum[i-1] + array[i]

    return prefix_sum

# Test the prefix sum function
my_array = [1, 2, 3, 4, 5]
prefix_sum = calculate_prefix_sum(my_array)
print(prefix_sum)

Explanation: In this example, we define a function calculate_prefix_sum that takes an array as input. This function calculates the prefix sum of the array.

We first create a new list called prefix_sum with the same length as the input array, initialized with zeros. This list will store the prefix sum values.

Next, we assign the first element of the input array to the corresponding index of prefix_sum, as the prefix sum of the first element is simply the element itself.

Then, we iterate over the remaining elements of the input array using a for loop, starting from index 1. For each element, we calculate the prefix sum by adding the current element to the prefix sum of the previous element. We update the prefix_sum list with the computed prefix sum.

Finally, we return the prefix_sum list containing the cumulative sums of the input array.

When we test the prefix sum function with my_array = [1, 2, 3, 4, 5], it calculates the prefix sum of these elements and stores the results in the prefix_sum list, which is then printed.

Now, let's solve a competitive question using the prefix sum technique. The question is "Subarray Sum Equals K," which requires finding the total number of continuous subarrays whose sum equals a given target value.

def subarray_sum(nums, k):
    prefix_sum = {0: 1}
    count = 0
    current_sum = 0

    for num in nums:
        current_sum += num
        complement = current_sum - k

        if complement in prefix_sum:
            count += prefix_sum[complement]

        prefix_sum[current_sum] = prefix_sum.get(current_sum, 0) + 1

    return count

# Example usage
nums = [1, 1, 1]
k = 2
print(subarray_sum(nums, k))
# Output: 2

In this solution, we use the prefix sum technique to count subarrays with a sum equal to k. We maintain a prefix_sum dictionary to store the cumulative sum encountered so far and the frequency of each sum. We initialize the dictionary with {0: 1} to handle the case when the sum of the subarray from the beginning is equal to k.

As we iterate over the array nums, we update the current_sum by adding the current element. We then calculate the complement, which is the difference between the current_sum and the target k. If the complement exists in the prefix_sum dictionary, it means there is a subarray with a sum equal to k. We increment the count by the frequency of that complement in the dictionary.

Finally, we update the prefix_sum dictionary by adding the current_sum to it, along with its frequency. The solution has a time complexity of O(n) since we traverse the array once, and the space complexity is O(n) because the dictionary can store up to n unique sums.

Next Post Previous Post