Range queries - creating prefix sum of a 2D array

A range query is a type of search that returns a record of data based on a given range of values. This can be useful in a variety of situations, such as finding minimum value in a range or sum of all the values in a range.

In this article, we will focus on a situation where arrays are idempotent/static i.e. array values are never updated between the queries.

Sum queries

Sum queries on a static array can be answered in O(1) time after constructing a prefix sum array from original array. For example, A[i] in prefix sum array will contain sum(0, A[i]) from original array.

Once a prefix array is constructed, sum queries can be answered using following formula

        sum(a,b) = sum(0, b) - sum(0, a - 1)

//sampleStart
fun main() {
    // original array
    val arr = intArrayOf(1, 3, 4, 8, 6, 1, 4, 2)

    // create prefix sum
    val prefixSum = IntArray(arr.size)
    prefixSum[0] = arr[0]
    for (i in 1 until arr.size) {
        prefixSum[i] = prefixSum[i - 1] + arr[i]
    }

    println("arr       : ${arr.contentToString()}")
    println("prefix sum: ${prefixSum.contentToString()}")

    val range = 2..5
    // sum(a,b) = sum(0,b) - sum(0, a-1)
    val sum = prefixSum[range.last] - prefixSum[range.first.dec()]
    println("sum of range $range = $sum")
}
//sampleEnd

Time complexity of creating a prefix sum array is O(n) and auxiliary space complexity is also O(n) because an extra array of same size as original is needed to store prefix sum.

Two dimensional range sum

The same idea of prefix sum array can be extended to multi dimentional arrays, however calculating the prefix sum for multi dimensional arrays is a bit tricky. Consider the following example

Let’s see how this prefix sum was calculated. For first row and column (i,j) = (0,0) prefix values are calculated as we did before.

The value at A[0][j + 1] = A[0][j] + A[0][j+1], similarly the value at A[i+1][0] = A[i+1][0] + A[i-1][0].

For all other values of (i, j), prefix sum is calculated by taking a square of 4 values into account. For example, to calculate value of (1, 1) in our example, we take values (0,0), (0,1), (1,0) and (1,1).

As sum for cells (0,1) & (1,0) is already calculated (5 and 2 respectively), we can just sum up the values in these two cells and current cell to get the value of (1,1). But there is one problem, in the above example you can see the value 1 was added twice, to correct this, we will subtract this value from the total sum. The formula for calculating the sum is

    sum[i][j] = A[i][j] + A[i-1][j] + A[i][j-1] - A[i-1][j-1] where i > 0 and j > 0
             = current cell + top cell  + left cell  - top-left cell

Using this formula we can calculate the value for (1,1) as below

    sum[1][1] = A[1][1] + A[0][1] + A[1][0] - A[0][0]
              = 1 + 5 + 2 - 1 = 7

Similarly, you can calculate the value for (1,2) and so on. Try to calculate the missing value in the following diagram

A conscious reader might have guessed by now that same formula can be applied to first row and column if we return 0 value for left, top and topLeft positions when indices are 0. I will using the same formula to calculate value for each cell in the code example below

fun prefixSum2D(array: Array<IntArray>): Array<IntArray> {
    if (array.isEmpty()) return array
    val rowSize = array[0].size
    val prefixSum = Array(array.size) { IntArray(rowSize) }

    for (i in array.indices) {
        for (j in 0 until rowSize) {
            // initialize current value
            prefixSum[i][j] = array[i][j]
            // calculate top, left, topLeft values
            val top = if (i == 0) 0 else prefixSum[i - 1][j]
            val left = if (j == 0) 0 else prefixSum[i][j - 1]
            val topLeft = if (i > 0 && j > 0) {
                prefixSum[i - 1][j - 1]
            } else {
                0
            }
            // calculate sum
            prefixSum[i][j] += top.plus(left).minus(topLeft)
        }
    }
    return prefixSum
}

Calculating the sum

In the above diagram, let’s say we have to calculate sum of values under red square. We can take full rectangle, colored blue in the diagram and can subtract top & left rectangles (colored orange and gray respectively). Again, the problem we have here is the value in cell (0,0) is subtracted twice, so to fix this, after subtracting both top and left rectangles we will add this value back to the sum.

The values in different rectangles are marked as a, b, c and d. The formula for calculating the sum would be

    S(a) - S(b) - S(c) + s(d)

Total sum of the red square in our example would be

    28 - 2 - 10 + 1 = 17

What if we do not have those top and left rectangles to subtract. Well in that case, the value S(a) will contain the actual sum. See the following example

There are no rectangles to subtract, so S(a) contains the actual sum which is 7.

This formula will only work if sum area is rectangle

Complete example in kotlin

//sampleStart
fun main() {
    // click + to see full script
    val array = Array(4) {
        IntArray(4)
    }
    array[0] = intArrayOf(1, 4, 5, 6)
    array[1] = intArrayOf(1, 1, 2, 3)
    array[2] = intArrayOf(0, 5, 9, 8)
    array[3] = intArrayOf(1, 2, 3, 0)

    println("Array")
    array.forEach {
        println(it.contentToString())
    }
    println()
    val sumTable = prefixSum2D(array)
    println("Prefix Sum Table")
    sumTable.forEach {
        println(it.contentToString())
    }

    println()

    val top = Pair(1, 1)
    val right = Pair(1, 2)
    val left = Pair(2, 1)
    val bottom = Pair(2, 2)
    getSum(
        sumTable = sumTable,
        top = top,
        right = right,
        left = left,
        bottom = bottom
    ).also {
        println("sum of $top $right $left $bottom = $it")
    }
}
//sampleEnd

fun prefixSum2D(array: Array<IntArray>): Array<IntArray> {
    if (array.isEmpty()) return array
    val rowSize = array[0].size
    val prefixSum = Array(array.size) { IntArray(rowSize) }

    for (i in array.indices) {
        for (j in 0 until rowSize) {
            // initialize current value
            prefixSum[i][j] = array[i][j]
            // calculate top, left, topLeft values
            val top = if (i == 0) 0 else prefixSum[i - 1][j]
            val left = if (j == 0) 0 else prefixSum[i][j - 1]
            val topLeft = if (i > 0 && j > 0) {
                prefixSum[i - 1][j - 1]
            } else {
                0
            }
            // calculate sum
            prefixSum[i][j] += top.plus(left).minus(topLeft)
        }
    }
    return prefixSum
}

fun getSum(
    sumTable: Array<IntArray>,
    top: Pair<Int, Int>, right: Pair<Int, Int>,
    left: Pair<Int, Int>, bottom: Pair<Int, Int>,
): Int {
    if (sumTable.isEmpty()) return 0

    val columnSize = sumTable[0].size
    val rows = sumTable.size

    (columnSize.dec()..rows.dec()).also {
        checkInRange(it.first, it.last, top)
        checkInRange(it.first, it.last, right)
        checkInRange(it.first, it.last, left)
        checkInRange(it.first, it.last, bottom)
    }


    val topRect = if (right.first - 1 < 0) {
        0
    } else {
        sumTable[right.first.dec()][right.second]
    }

    val topLeftRect = if (top.first.dec() < 0 || top.second.dec() < 0) {
        0
    } else {
        sumTable[top.first.dec()][top.second.dec()]
    }

    val leftRect = if (left.second - 1 < 0) {
        0
    } else {
        sumTable[left.first][left.second.dec()]
    }

    val bottomRect = sumTable[bottom.first][bottom.second]
    return bottomRect - leftRect - topRect + topLeftRect
}

fun checkInRange(column: Int, rows: Int, range: Pair<Int, Int>) {
    if (range.first < 0 || range.first > rows) {
        error("invalid range $range")
    }

    if (range.second < 0 || range.second > column) {
        error("invalid range $range")
    }
}