Problem Description

You are given an array nums of length n and a positive integer k.

A subarray is considered good if the absolute difference between its first element and last element equals exactly k. In other words, for a subarray nums[i..j], it is good when |nums[i] - nums[j]| == k.

Your task is to find the maximum sum among all good subarrays in the array. If no good subarrays exist, return 0.

For example, if nums = [1, 2, 3, 4, 5] and k = 3, a subarray [2, 3, 4, 5] would be good because |2 - 5| = 3, which equals k. The sum of this subarray would be 2 + 3 + 4 + 5 = 14.

The challenge involves:

• Finding all subarrays where the absolute difference between the first and last elements equals k
• Calculating the sum of each such subarray
• Returning the maximum sum found, or 0 if no valid subarrays exist

Intuition

To find the maximum sum of a good subarray, we need to efficiently identify all subarrays where |nums[i] - nums[j]| == k.

Let's think about this problem differently. For any element nums[j] at position j, we want to find if there exists a previous element nums[i] (where i < j) such that:

• Either nums[i] = nums[j] - k
• Or nums[i] = nums[j] + k

If we find such an element, then the subarray from position i to j is a good subarray.

Now, to calculate the sum of subarray nums[i..j], we can use prefix sums. If we know the prefix sum up to position j (let's call it sum[0..j]) and the prefix sum up to position i-1 (let's call it sum[0..i-1]), then the subarray sum is simply sum[0..j] - sum[0..i-1].

Here's the key insight: for each position j, we want to find the position i where nums[i] equals either nums[j] - k or nums[j] + k, and the prefix sum sum[0..i-1] is as small as possible (to maximize the subarray sum).

This leads us to use a hash table that maps each value to the minimum prefix sum seen before that value appears. As we iterate through the array:

1. We check if nums[j] - k or nums[j] + k exists in our hash table
2. If yes, we calculate the subarray sum and update our maximum
3. We update the hash table with the current element and its corresponding prefix sum (only if it gives a smaller prefix sum for that value)

The reason we store the minimum prefix sum for each value is that when we later find a matching element to form a good subarray, subtracting a smaller prefix sum will give us a larger subarray sum, which is what we want to maximize.

Learn more about Prefix Sum patterns.

Solution Approach

We implement the solution using prefix sums and a hash table to efficiently find and calculate the maximum sum of good subarrays.

Data Structures:

• Hash table p: Maps each element value to the minimum prefix sum seen before that element appears
• Variable s: Tracks the current prefix sum as we iterate through the array
• Variable ans: Stores the maximum sum found, initialized to negative infinity

Algorithm Steps:

1. Initialize the hash table: Set p[nums[0]] = 0, which represents that before nums[0], the prefix sum is 0. This allows us to consider subarrays that start from index 0.

2. Iterate through the array: For each element nums[i] with its value x:

   • Update the prefix sum: s += x (now s represents the sum from nums[0] to nums[i])

3. Check for good subarrays ending at position i:

   • If x - k exists in p: We found an element earlier that can form a good subarray with the current element. The subarray sum is s - p[x - k]. Update ans = max(ans, s - p[x - k]).
   • If x + k exists in p: Similarly, this forms another good subarray. Update ans = max(ans, s - p[x + k]).

4. Update the hash table for future iterations:

   • We prepare for the next element nums[i + 1] (if it exists)
   • Only update p[nums[i + 1]] if:
     • nums[i + 1] is not in p (first occurrence)
     • OR p[nums[i + 1]] > s (found a smaller prefix sum for the same value)
   • This ensures we always keep the minimum prefix sum for each value

5. Return the result:

   • If ans is still negative infinity, no good subarray was found, return 0
   • Otherwise, return ans

Why this works:

• When we find x - k or x + k in the hash table, it means there's a previous position where that value appeared
• The prefix sum stored in p for that value represents the sum before that position
• By subtracting this from the current prefix sum s, we get the exact sum of the subarray between those two positions
• By always storing the minimum prefix sum for each value, we maximize the subarray sum when we calculate s - p[value]

Example Walkthrough

Let's walk through the solution with nums = [1, 5, 3, 4, 2] and k = 3.

We need to find subarrays where |first - last| = 3 and return the maximum sum.

Initial Setup:

• p = {1: 0} (before nums[0]=1, prefix sum is 0)
• s = 0 (current prefix sum)
• ans = -∞ (maximum sum found)

Iteration 1: i=0, x=nums[0]=1

• Update prefix sum: s = 0 + 1 = 1
• Check for good subarrays:
  • x - k = 1 - 3 = -2: Not in p
  • x + k = 1 + 3 = 4: Not in p
• Update hash table for next element:
  • Next element is nums[1]=5, not in p
  • Set p[5] = 1 (prefix sum before position 1)
• Current state: p = {1: 0, 5: 1}, ans = -∞

Iteration 2: i=1, x=nums[1]=5

• Update prefix sum: s = 1 + 5 = 6
• Check for good subarrays:
  • x - k = 5 - 3 = 2: Not in p
  • x + k = 5 + 3 = 8: Not in p
• Update hash table for next element:
  • Next element is nums[2]=3, not in p
  • Set p[3] = 6 (prefix sum before position 2)
• Current state: p = {1: 0, 5: 1, 3: 6}, ans = -∞

Iteration 3: i=2, x=nums[2]=3

• Update prefix sum: s = 6 + 3 = 9
• Check for good subarrays:
  • x - k = 3 - 3 = 0: Not in p
  • x + k = 3 + 3 = 6: Not in p
• Update hash table for next element:
  • Next element is nums[3]=4, not in p
  • Set p[4] = 9 (prefix sum before position 3)
• Current state: p = {1: 0, 5: 1, 3: 6, 4: 9}, ans = -∞

Iteration 4: i=3, x=nums[3]=4

• Update prefix sum: s = 9 + 4 = 13
• Check for good subarrays:
  • x - k = 4 - 3 = 1: Found in p! p[1] = 0
    • Subarray sum = s - p[1] = 13 - 0 = 13
    • This represents subarray [1,5,3,4] where |1-4| = 3 ✓
    • Update ans = max(-∞, 13) = 13
  • x + k = 4 + 3 = 7: Not in p
• Update hash table for next element:
  • Next element is nums[4]=2, not in p
  • Set p[2] = 13 (prefix sum before position 4)
• Current state: p = {1: 0, 5: 1, 3: 6, 4: 9, 2: 13}, ans = 13

Iteration 5: i=4, x=nums[4]=2

• Update prefix sum: s = 13 + 2 = 15
• Check for good subarrays:
  • x - k = 2 - 3 = -1: Not in p
  • x + k = 2 + 3 = 5: Found in p! p[5] = 1
    • Subarray sum = s - p[5] = 15 - 1 = 14
    • This represents subarray [5,3,4,2] where |5-2| = 3 ✓
    • Update ans = max(13, 14) = 14
• No next element to update
• Final result: ans = 14

The maximum sum good subarray is [5,3,4,2] with sum 14.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm iterates through the array exactly once using a single for loop. Within each iteration:

• Dictionary lookups (x - k in p and x + k in p) take O(1) time on average
• Dictionary insertions and updates (p[nums[i + 1]] = s) take O(1) time on average
• All other operations (comparisons, arithmetic operations, max calculations) take O(1) time

Since we perform n iterations with O(1) operations per iteration, the overall time complexity is O(n).

Space Complexity: O(n)

The space usage comes from:

• Dictionary p which stores prefix sums indexed by array elements. In the worst case, when all elements are unique, the dictionary can store up to n key-value pairs
• Variable storage (ans, s, n, i, x) uses O(1) space

Therefore, the space complexity is dominated by the dictionary, resulting in O(n) space complexity, where n is the length of the input array.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Incorrect Hash Table Update Timing

The Problem: A critical mistake is updating the hash table with the current element's prefix sum BEFORE checking for good subarrays. This can lead to considering single-element subarrays when k = 0.

Incorrect Implementation:

for i, current_num in enumerate(nums):
    prefix_sum += current_num
  
    # WRONG: Updating before checking
    if current_num not in min_prefix_before:
        min_prefix_before[current_num] = prefix_sum - current_num
  
    # Now checking - but hash table already has current element!
    if current_num - k in min_prefix_before:
        max_sum = max(max_sum, prefix_sum - min_prefix_before[current_num - k])

When k = 0, the condition current_num - k in min_prefix_before becomes current_num in min_prefix_before, which is now always true after our update. This incorrectly considers single-element "subarrays" as valid.

The Solution: Always update the hash table for the NEXT element, not the current one. This ensures we only consider subarrays with at least 2 elements:

# Correct: Update for next element after checking current
if i + 1 < n:
    next_num = nums[i + 1]
    if next_num not in min_prefix_before or min_prefix_before[next_num] > prefix_sum:
        min_prefix_before[next_num] = prefix_sum

Pitfall 2: Not Maintaining Minimum Prefix Sum

The Problem: Simply storing any prefix sum for each value, rather than the minimum, leads to suboptimal results.

Incorrect Implementation:

# WRONG: Always overwriting, not keeping minimum
if i + 1 < n:
    min_prefix_before[nums[i + 1]] = prefix_sum

Why it fails: Consider nums = [1, 5, 1, 3] with k = 4. If we don't maintain the minimum:

• After processing index 0: min_prefix_before[5] = 1
• After processing index 1: min_prefix_before[1] = 6 (overwriting)
• At index 2 (value=1), checking for 1 + 4 = 5: we get sum = 7 - 6 = 1

But if we maintained the minimum:

• min_prefix_before[1] would still be 0 (from initialization)
• At index 2, we'd get sum = 7 - 1 = 6 (correct maximum)

The Solution: Only update when finding a smaller prefix sum:

if next_num not in min_prefix_before or min_prefix_before[next_num] > prefix_sum:
    min_prefix_before[next_num] = prefix_sum

Pitfall 3: Forgetting Edge Cases

The Problem: Not properly handling when no good subarrays exist, leading to returning incorrect values like negative infinity or causing errors.

The Solution: Always check if max_sum was updated from its initial value before returning:

return 0 if max_sum == float('-inf') else max_sum