Thoughts of a programmer

Square Root Decomposition

Let’s consider the Problem Statement from HackerEarth.

Basically, you are given an array \(A\) of size \(N\) and \(Q\) queries. Each query can be of two types -

And all the indexes are \(1\)-based.

So how can we solve this problem. Let’s try bruteforce approach. Here is a sample implementation in Python 3.

n, q = list(map(int, input()s.strip().split(' ')))
array = list(map(int, input().strip().split(' ')))
for i in range(q):
    typ, k = list(map(int, input().strip().split(' ')))
    if typ == 1:
        array[k - 1] += 1
    else:
        odd_count = 0
        found_at = -1
        for index, num in enumerate(array):
            if num & 1 > 0:
                odd_count += 1
                if odd_count == k:
                    found_at = index + 1
                    break
        print(found_at)

Here we are simply iterating over the queries. And depending on whether the query is of type \(1\) or type \(2\), we either update the corresponding \(i^{th}\) entry in the array or loop over all the elements to find the \(K^{th}\) odd number.

But this approach would only work when the input size is small. For this problem, the input array could be as large as \(2 * 10 ^ 6\) and the number of queries could also be as large as \(2 * 10 ^ 6\). In the worst case (that would occur when the \(K^{th}\) odd number doesn’t exists and all the queries are of type \(2\)), this would mean iterating over whole of the array of size \(2 * 10 ^ 6\), \(2 * 10 ^ 6\) times. Which is clearly not gonna cut it.

Thankfully, we can solve this problem under the given constraints with a neat little technique called Square root decomposition.

Idea

Main idea behind the Square root decomposition method is to divide our original array into buckets of equal size, usually \(\sqrt{n}\) where \(n\) is the size of the original array.

Dividing the array into the buckets allows us to precompute some information about those buckets and later reuse them when given some query over the array.

For example, suppose we have to find some information about a range \(l\) to \(r\). Now the given query range would be covering some of the buckets partially and others completely. Now the thing to note is that there could only be at most \(2\) partially covered buckets. One on the left and one on the right of the query range.

For all the completely covered buckets, we won’t have to compute anything. We could directly use the precomputed information from above.

Only for these \(2\) partially covered buckets, we would have to go inside those buckets, and compute the necessary information.

Finally, we can merge the collected information from these buckets to answer the given query.

In case we need to update the value of some array element, we would only need to recompute the information for the corresponding bucket only. And that could be done in \(O(\sqrt{N})\) time.

Solution

Now having gone through the basic idea of the Square Root Decomposition, let’s see how we can use this to solve our original problem statement.

First we will divide the original array into \(\sqrt{n}\) buckets. Here \(n\) is at most \(2 * 10 ^ 6\), so number of buckets could be atmost \(\lceil\sqrt{2 * 10 ^ 6}\rceil=1415\) and there would be \(\lceil\frac{2 * 10 ^ 6}{1415}\rceil=1414\) numbers in each bucket with the last bucket having just \(2 * 10 ^ 6 mod 1414=604\) elements.

2
2
1
1
4
4
5
5
6
6
7
7
3
3
11
11
67
67
43
43
1
1
2
2
1
1
3
3
Original array
[Not supported by viewer]
Decomposed Odd Count
[Not supported by viewer]
n = 10
buckets = ceil(sqrt(10)) = 4
bucket size = ceil(10 / 4) = 3
[Not supported by viewer]

Now for each bucket, we can compute the count of odd numbers in that bucket. As there can be maximum of \(1415\) buckets, we will have \(1415\) counts of odd numbers.

Now whenever we will receive the query of type 1, we can update the corresponding number in the original array and depending upon whether the new number is odd or even, we can either increase or decrease the odd number count of the corresponding bucket. This operation can be done \(O(1)\) time.

When we receive a query of type 2, instead of iterating over whole array like we were doing in the brute force approach, we can instead take use of our odd number count that we precalculated.

Suppose we need to find the \(100^{th}\) odd number and assume that it exists at position \(10 ^ 6\). We will iterate over our decomposed counter array and maintain a cumulative sum of the total odd numbers seen till now. When we have seen at least \(100\) odd numbers, we can then break and loop for that particular bucket which caused the sum to exceed \(100\) and find the corresponding element position. This way we will just have to iterate for \(\lceil\frac{10 ^ 6}{1414}\rceil=708\) buckets, then over atmost \(1414\) elements inside that last bucket. Compare that to the \(10 ^ 6\) lookups done by the brute force approach.

Time complexity for this operation will be thus \(O(NumBuckets + BucketSize)\). Since number of buckets and bucket size are almost the same, we can say that complexity is \(O(NumBuckets)=O(\sqrt{n})\).

2
2
1
1
4
4
5
5
6
6
7
7
3
3
11
11
67
67
43
43
1
1
2
2
1
1
3
3
Finding  5th Odd Number
[Not supported by viewer]
Count 6
Count 6
Count 1
Count 1
Count 3
Count 3
Count exceeds 5
Count exceeds 5
Count matches 5
Count matches 5

Overall complexity of the algorithm will be thus the time for precomputation of the buckets odd counts plus the time to answer the queries. The complexity then is \(O(n + q\sqrt{n})\).

Here is a sample implementation of the above strategy in Python 3.

import math
n, q = list(map(int, input().strip().split(' ')))
array = list(map(int, input().strip().split(' ')))
# Taking ceil, to cover the last bucket also
buckets = int(math.ceil(math.sqrt(n)))
bucket_size = int(math.ceil(n / buckets))
odd_count = [0 for _ in range(buckets)]
# Precalculate buckets odd number count
for bucket in range(buckets):
    start = bucket * bucket_size
    # Taking into fact that last bucket might have less number of elements than bucket_size
    end = min(n, start + bucket_size)
    for index in range(start, end):
        if array[index] & 1 > 0:
            odd_count[bucket] += 1
for i in range(q):
    typ, k = list(map(int, input().strip().split(' ')))
    if typ == 1:
        k -= 1
        array[k] += 1
        # Update the odd count for the corresponding bucket
        bucket = k // bucket_size
        odd_count[bucket] += 1 if (array[k] & 1 > 0) else -1
    else:
        cnt = 0
        found_at = -1
        for index, count in enumerate(odd_count):
            cnt += count
            if cnt >= k:
                # Then dig deeper into the corresponding bucket to find the kth element
                bucket = index
                start = bucket * bucket_size
                end = min(n, start + bucket_size)
                # Decrease the count to count again for the current bucket
                cnt -= count
                for j in range(start, end):
                    if array[j] & 1 > 0:
                        cnt += 1
                        if cnt == k:
                            found_at = j + 1 # 1-based indexing
                            break
                break
        print(found_at)

You can find more problems related to Square root decomposition here.

Thank you for reading my article. Let me know if you liked my article or any other suggestions for me, in the comments section below. And please, feel free to share :)