Advanced Prefix Sum: 2D, Difference Arrays, and Beyond

Prefix sums are one of the most versatile tools in competitive programming and interviews. Once you master the 1D version, a whole world opens up — 2D prefix sums for matrix queries, difference arrays for batch range updates, and modular-arithmetic tricks that turn brute-force O(n^2) solutions into elegant O(n) ones.

---

1. 2D Prefix Sum — Submatrix Sum Queries

Given a matrix of integers, you often need to answer many queries of the form: “What is the sum of all elements in the submatrix from (r1, c1) to (r2, c2)?” A naive approach takes O(rows x cols) per query. With a 2D prefix sum, each query takes O(1).

Building the 2D Prefix Sum

The idea extends naturally from 1D. Define prefix[i][j] as the sum of all elements in the submatrix from (0, 0) to (i-1, j-1).

def build_2d_prefix(matrix):
    """
    Build a 2D prefix sum matrix.

    prefix[i][j] = sum of matrix[0..i-1][0..j-1]

    Time: O(m * n), Space: O(m * n)
    """
    if not matrix or not matrix[0]:
        return [[]]

    m, n = len(matrix), len(matrix[0])
    prefix = [[0] * (n + 1) for _ in range(m + 1)]

    for i in range(1, m + 1):
        for j in range(1, n + 1):
            prefix[i][j] = (
                matrix[i - 1][j - 1]
                + prefix[i - 1][j]
                + prefix[i][j - 1]
                - prefix[i - 1][j - 1]
            )

    return prefix

Querying a Submatrix Sum

To get the sum of the submatrix from (r1, c1) to (r2, c2) (0-indexed), we use inclusion-exclusion:

def submatrix_sum(prefix, r1, c1, r2, c2):
    """
    Query the sum of submatrix [r1..r2][c1..c2] in O(1).

    Uses inclusion-exclusion on the 2D prefix sum.
    """
    return (
        prefix[r2 + 1][c2 + 1]
        - prefix[r1][c2 + 1]
        - prefix[r2 + 1][c1]
        + prefix[r1][c1]
    )

Complexity

Operation	Time	Space
Build prefix	O(m * n)	O(m * n)
Query submatrix sum	O(1)	O(1)
Q queries total	O(m * n + Q)	O(m * n)

Example Walkthrough

matrix = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
]

prefix = build_2d_prefix(matrix)

# Sum of submatrix (1,1) to (2,2) = 5 + 6 + 8 + 9 = 28
print(submatrix_sum(prefix, 1, 1, 2, 2))  # Output: 28

# Sum of entire matrix (0,0) to (2,2) = 45
print(submatrix_sum(prefix, 0, 0, 2, 2))  # Output: 45

# Single element (0,0) to (0,0) = 1
print(submatrix_sum(prefix, 0, 0, 0, 0))  # Output: 1

LeetCode: 304. Range Sum Query 2D - Immutable

class NumMatrix:
    def __init__(self, matrix):
        m, n = len(matrix), len(matrix[0])
        self.prefix = [[0] * (n + 1) for _ in range(m + 1)]

        for i in range(1, m + 1):
            for j in range(1, n + 1):
                self.prefix[i][j] = (
                    matrix[i - 1][j - 1]
                    + self.prefix[i - 1][j]
                    + self.prefix[i][j - 1]
                    - self.prefix[i - 1][j - 1]
                )

    def sumRegion(self, r1, c1, r2, c2):
        return (
            self.prefix[r2 + 1][c2 + 1]
            - self.prefix[r1][c2 + 1]
            - self.prefix[r2 + 1][c1]
            + self.prefix[r1][c1]
        )

Time: O(1) per query, O(m * n) initialization. Space: O(m * n).

---

2. Difference Arrays — Range Updates in O(1)

A difference array is the inverse of a prefix sum. Instead of answering range queries, it lets you perform range updates efficiently. If you need to add a value val to every element in the range [l, r], a difference array does it in O(1) per update.

The Idea

Given an array arr of length n, the difference array diff is defined so that prefix_sum(diff) = arr. To add val to all elements in [l, r]:

1. diff[l] += val
2. diff[r + 1] -= val (if r + 1 {'<'} n)

After all updates, take the prefix sum of diff to get the final array.

def apply_range_updates(n, updates):
    """
    Apply multiple range updates [l, r, val] to an array of zeros.

    Each update adds val to all positions in [l, r].

    Time: O(n + k) where k = number of updates
    Space: O(n)
    """
    diff = [0] * (n + 1)

    for l, r, val in updates:
        diff[l] += val
        if r + 1 <= n:
            diff[r + 1] -= val

    # Build result by taking prefix sum of diff
    result = [0] * n
    current = 0
    for i in range(n):
        current += diff[i]
        result[i] = current

    return result

Example

n = 5
updates = [
    (1, 3, 2),   # Add 2 to indices 1..3
    (2, 4, 3),   # Add 3 to indices 2..4
    (0, 2, -1),  # Subtract 1 from indices 0..2
]

result = apply_range_updates(n, updates)
print(result)  # [-1, 1, 4, 5, 3]

# Manual verification:
# Index:    0    1    2    3    4
# Start:    0    0    0    0    0
# +2 [1,3]: 0    2    2    2    0
# +3 [2,4]: 0    2    5    5    3
# -1 [0,2]:-1    1    4    5    3  ✓

Complexity

Operation	Time	Space
Single range update	O(1)	—
k updates + build result	O(n + k)	O(n)
Naive approach (k updates)	O(n * k)	O(n)

2D Difference Array

The same idea extends to 2D. To add val to all cells in the submatrix (r1, c1) to (r2, c2):

def apply_2d_range_updates(m, n, updates):
    """
    Apply submatrix range updates on an m x n grid.

    Each update: (r1, c1, r2, c2, val)

    Time: O(m * n + k), Space: O(m * n)
    """
    diff = [[0] * (n + 2) for _ in range(m + 2)]

    for r1, c1, r2, c2, val in updates:
        diff[r1][c1] += val
        diff[r1][c2 + 1] -= val
        diff[r2 + 1][c1] -= val
        diff[r2 + 1][c2 + 1] += val

    # Take 2D prefix sum to get the result
    result = [[0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            diff[i][j] += (
                (diff[i - 1][j] if i > 0 else 0)
                + (diff[i][j - 1] if j > 0 else 0)
                - (diff[i - 1][j - 1] if i > 0 and j > 0 else 0)
            )
            result[i][j] = diff[i][j]

    return result

---

3. Subarray Sum Divisible by K

Problem: Given an array of integers and an integer K, find the count of subarrays whose sum is divisible by K.

Key Insight: Modular Arithmetic + Pigeonhole

If prefix[j] % K == prefix[i] % K, then sum(arr[i..j-1]) % K == 0. So we count how many prefix sums share the same remainder when divided by K.

def subarrays_divisible_by_k(nums, k):
    """
    Count subarrays whose sum is divisible by k.

    Uses prefix sum + modular arithmetic.
    For each remainder r, if c prefix sums have remainder r,
    then C(c, 2) = c*(c-1)/2 subarrays exist.

    Time: O(n), Space: O(k)
    """
    remainder_count = [0] * k
    remainder_count[0] = 1  # empty prefix has sum 0
    prefix_sum = 0
    count = 0

    for num in nums:
        prefix_sum += num
        remainder = prefix_sum % k
        # Python's % handles negatives correctly (always non-negative)
        count += remainder_count[remainder]
        remainder_count[remainder] += 1

    return count

Example

nums = [4, 5, 0, -2, -3, 1]
k = 5

print(subarrays_divisible_by_k(nums, k))  # Output: 7

# The 7 subarrays:
# [4, 5, 0, -2, -3, 1] sum=5
# [5]                   sum=5
# [5, 0]                sum=5
# [5, 0, -2, -3]        sum=0
# [0]                   sum=0
# [-2, -3]              sum=-5
# [4, 5, 0, -2, -3]     sum=4+5+0-2-3=4... 
# Actually [0, -2, -3] sum=-5, etc.

Time: O(n). Space: O(k).

---

4. Continuous Subarray Sum (LeetCode 523)

Problem: Given an array and integer k, check if there exists a continuous subarray of size at least 2 whose sum is a multiple of k.

Pigeonhole Principle

By the pigeonhole principle, if the array has more than k elements, two prefix sums must share the same remainder mod k, guaranteeing a valid subarray. For smaller arrays, use a hash map.

def check_subarray_sum(nums, k):
    """
    Check if a continuous subarray of size >= 2 has sum
    that is a multiple of k.

    Store the first index where each remainder appears.
    If the same remainder appears again at index j where
    j - stored_index >= 2, we found a valid subarray.

    Time: O(n), Space: O(min(n, k))
    """
    # remainder -> first index where this remainder was seen
    remainder_index = {0: -1}  # empty prefix at index -1
    prefix_sum = 0

    for i, num in enumerate(nums):
        prefix_sum += num
        remainder = prefix_sum % k

        if remainder in remainder_index:
            if i - remainder_index[remainder] >= 2:
                return True
            # Don't update — we want the earliest index
        else:
            remainder_index[remainder] = i

    return False

Example

print(check_subarray_sum([23, 2, 4, 6, 7], 6))  # True: [2, 4] sums to 6
print(check_subarray_sum([23, 2, 6, 4, 7], 6))  # True: [23, 2, 6, 4, 7] sums to 42
print(check_subarray_sum([23, 2, 6, 4, 7], 13)) # False

Time: O(n). Space: O(min(n, k)).

---

5. Count Subarrays with Sum Equal to Zero

Problem: Count the number of subarrays whose sum equals zero.

This is a direct application of the prefix sum + hash map pattern. If prefix[i] == prefix[j], then sum(arr[i..j-1]) == 0.

def count_zero_sum_subarrays(nums):
    """
    Count subarrays with sum exactly 0.

    If a prefix sum value repeats c times, it contributes
    C(c, 2) = c*(c-1)/2 zero-sum subarrays.

    Time: O(n), Space: O(n)
    """
    from collections import defaultdict

    prefix_count = defaultdict(int)
    prefix_count[0] = 1
    prefix_sum = 0
    count = 0

    for num in nums:
        prefix_sum += num
        count += prefix_count[prefix_sum]
        prefix_count[prefix_sum] += 1

    return count

Example

print(count_zero_sum_subarrays([1, -1, 2, -2]))  # 3
# Subarrays: [1, -1], [2, -2], [1, -1, 2, -2]

print(count_zero_sum_subarrays([0, 0, 0]))  # 6
# All subarrays of [0, 0, 0] sum to 0:
# [0], [0], [0], [0,0], [0,0], [0,0,0]

print(count_zero_sum_subarrays([3, 4, -7, 3, 1, 3, 1, -4, -2, -2]))  # 4

Time: O(n). Space: O(n).

---

6. XOR Prefix — Range XOR Queries

XOR has a beautiful property: a ^ a = 0 and a ^ 0 = a. This means XOR prefix sums work just like regular prefix sums for answering range queries.

def build_xor_prefix(arr):
    """
    Build a prefix XOR array.

    prefix[i] = arr[0] ^ arr[1] ^ ... ^ arr[i-1]

    Time: O(n), Space: O(n)
    """
    n = len(arr)
    prefix = [0] * (n + 1)
    for i in range(n):
        prefix[i + 1] = prefix[i] ^ arr[i]
    return prefix


def range_xor(prefix, l, r):
    """
    XOR of elements from index l to r (inclusive) in O(1).
    """
    return prefix[r + 1] ^ prefix[l]

Applications of XOR Prefix

Find the Missing Number

def find_missing(nums, n):
    """
    nums contains n-1 distinct integers from [0, n].
    Find the missing one.

    XOR all numbers 0..n with all elements in nums.
    Time: O(n), Space: O(1)
    """
    xor_all = 0
    for i in range(n + 1):
        xor_all ^= i
    for num in nums:
        xor_all ^= num
    return xor_all

Count Subarrays with XOR Equal to Target

def count_xor_subarrays(nums, target):
    """
    Count subarrays whose XOR equals target.

    Similar to subarray sum = k but with XOR.
    prefix[j] ^ prefix[i] = target
    => prefix[i] = prefix[j] ^ target

    Time: O(n), Space: O(n)
    """
    from collections import defaultdict

    prefix_count = defaultdict(int)
    prefix_count[0] = 1
    xor_so_far = 0
    count = 0

    for num in nums:
        xor_so_far ^= num
        # We need prefix_count[xor_so_far ^ target]
        count += prefix_count[xor_so_far ^ target]
        prefix_count[xor_so_far] += 1

    return count

Example

arr = [4, 2, 2, 6, 4]

prefix = build_xor_prefix(arr)
print(range_xor(prefix, 1, 3))  # 2 ^ 2 ^ 6 = 6

print(count_xor_subarrays([4, 2, 2, 6, 4], 6))  # 4
# Subarrays with XOR=6: [4,2], [2,2,6], [6], [2,2,6,4,... no]

---

7. Summary of Techniques

Technique	Use Case	Time	Space
2D Prefix Sum	Submatrix sum queries	O(1) per query	O(m * n)
Difference Array	Batch range updates	O(1) per update	O(n)
Prefix % K	Subarrays divisible by K	O(n)	O(k)
Prefix + Hash Map	Subarrays with sum = 0	O(n)	O(n)
XOR Prefix	Range XOR queries	O(1) per query	O(n)
XOR + Hash Map	Subarrays with XOR = target	O(n)	O(n)

When to Recognize These Patterns

• “Sum of submatrix” or “rectangle sum” → 2D prefix sum
• “Add val to range [l, r]” with many updates → difference array
• “Divisible by K” → prefix sum mod K + counting remainders
• “Subarray sum equals 0” → prefix sum + hash map for repeats
• “XOR of a range” → XOR prefix (since XOR is its own inverse)

---

8. Practice Problems

Problem	Platform	Difficulty	Key Technique
Range Sum Query 2D	LeetCode 304	Medium	2D prefix sum
Matrix Block Sum	LeetCode 1314	Medium	2D prefix sum
Subarray Sums Divisible by K	LeetCode 974	Medium	Prefix % K
Continuous Subarray Sum	LeetCode 523	Medium	Prefix % K + index
Subarray Sum Equals K	LeetCode 560	Medium	Prefix + hash map
Count Number of Nice Subarrays	LeetCode 1248	Medium	Prefix sum variant
XOR Queries of a Subarray	LeetCode 1310	Medium	XOR prefix
Corporate Flight Bookings	LeetCode 1109	Medium	Difference array
Car Pooling	LeetCode 1094	Medium	Difference array
Range Addition	LeetCode 370	Medium	Difference array

---

Key Takeaways

1. 2D prefix sums extend naturally from 1D — build with inclusion-exclusion, query with inclusion-exclusion.
2. Difference arrays are the inverse of prefix sums — O(1) range updates, then one prefix sum pass to materialize.
3. Modular arithmetic with prefix sums unlocks divisibility problems — count remainders, not sums.
4. XOR prefix works exactly like sum prefix because XOR is associative, commutative, and self-inverse.
5. These techniques combine: a 2D difference array handles batch submatrix updates; XOR prefix + hash map solves “count subarrays with XOR = target.”