mo algorithm for summing the subarray

Question

Below is the implementation of MO algorithm for summing the sub array given the range as input. As MO algorithm is a offline query it is given a list of ranges as input.

Below code calculates the result using MO algorithm and compares the result with expected output by directly summing the sub array.

MO algorithm first divides the whole array in sqrt(len of array) blocks and then divides the queries based on blocks. It sorts the input queries based on blocks and then based on the right parameters if the block number is same. It uses two pointers which moves left and right based on input query. The main intention of the algorithm is to reduce the movements of the pointers.

import collections
import math

def cmp(a, b):
    if a > b:
        return 1
    elif a < b:
        return -1
    return 0

def compare(a, b):
    def bin(x):
        return x[0]/r
    result = cmp(a, b)
    if result:
        return result
    return cmp(x[1], y[1])

a   = [1, 1, 2, 1, 3, 4, 5, 2, 8];
r = math.sqrt(len(a))
d_query = {0:[0, 5], 1:[5, 8], 2:[2, 4], 3:[5, 6]}
od = collections.OrderedDict(sorted(d_query.items(), compare, key= lambda x: x[1]))
print(od)

l = r = 0
cur_sum = 0
cur_l = 0
cur_r = -1
result = {}
for key, value in od.items():
    l = value[0]
    r = value[1]
    while cur_r < r:
        cur_r += 1
        cur_sum += a[cur_r]
    while cur_r > r:
        cur_sum -= a[cur_r]
        cur_r -= 1
    while cur_l < l:
        cur_sum -= a[cur_l]
        cur_l += 1
    while cur_l > l:
        cur_l -= 1
        cur_sum += a[cur_l]
    result[key] = cur_sum

for (k1,v1), (k2, v2) in zip(d_query.items(),result.items()):
    if k1 != k2 or v1 == v2:
        print("failed")
        break
print(result)

Answer 1

• I would use functools.partial() to pass the block size into the comparator
• choose more explicit and readable variable names - left, right and block_size instead of l, r and r

• there is a simpler way to create an expected result dictionary:

  expected_result = {key: sum(a[left:right+1]) for key, (left, right) in queries.items()}
  

• we can unpack left and right range borders instead of getting them by index
• note that there is no cmp in Python 3 anymore - you would have to use the cmp_to_key() helper to make it work

Improved code:

from functools import partial
import math


def compare(block_size, x, y):
    """The heart of the Mo's algorithm - determines if a range preceeds another."""
    _, (x_left, x_right) = x
    _, (y_left, y_right) = y

    if x_left / block_size != y_left / block_size:
        return x_left / block_size < y_left / block_size

    return x_right < y_right


def get_expected_result(a, queries):
    """Calculates given queries as sums of given ranges."""
    return {key: sum(a[left:right + 1]) for key, (left, right) in queries.items()}


def mos_algorithm(a, queries):
    """Mo's algorithm implementation."""
    block_size = math.sqrt(len(a))

    cur_sum, cur_left, cur_right = 0, 0, -1
    result = {}
    for key, (left, right) in sorted(queries.items(), cmp=partial(compare, block_size)):
        while cur_right < right:
            cur_right += 1
            cur_sum += a[cur_right]

        while cur_right > right:
            cur_sum -= a[cur_right]
            cur_right -= 1

        while cur_left < left:
            cur_sum -= a[cur_left]
            cur_left += 1

        while cur_left > left:
            cur_left -= 1
            cur_sum += a[cur_left]

        result[key] = cur_sum

    return result


if __name__ == '__main__':
    a = [1, 1, 2, 1, 3, 4, 5, 2, 8]
    queries = {0: [0, 5], 1: [5, 8], 2: [2, 4], 3: [5, 6]}

    result = mos_algorithm(a, queries)
    expected_result = get_expected_result(a, queries)

    print(result)
    print(expected_result)

Prints:

{0: 12, 1: 19, 2: 6, 3: 9}
{0: 12, 1: 19, 2: 6, 3: 9}