Count divisible elements in subarray

Question

I am trying to solve the following exercise:

There is an array of positive integers with size N given 1 <= N <= 2*10⁵. Each element of the array is at most 10⁵ at all times.
After reading in these N integers, you need to support Q queries (1 <= Q <= 10⁵).

Each query asks one of the below:

1. Count divisibility (represented by letter Q in input):
   Find the number of elements in a subarray that are divisible by a given input value. For this query, you are given 3 integers: the left end of the subarray, the right end of the subarray, and the integer to count divisibility for.

2. Update (represented by letter U in input):
   Change the element at some index to a new value. For this query, you are given 2 integers: the index to change and the new value.

I have a O(NQ) brute force solution but it times out. Any ideas how to solve this in better time complexity?

---

The input is N and Q on the first line. The next line has all the starting elements of the array separated by spaces.

The next Q lines each contain one query. The parts of the query are separated by spaces and the indexes are 0 based.

Example input:

6 4
6 5 4 3 2 1
Q 0 3 1
Q 1 4 2
U 3 8
Q 1 5 2

The ideal output is (each query for the count of divisible elements should be answered on a new line)

4
2
3

My brute force code is this (too slow however):

#include <iostream>
#include <vector>

int main() {
    using std::cin, std::cout;

    int N, Q;
    cin >> N >> Q;

    std::vector<int> array(N);
    for (int &i : array) cin >> i; // read starting array elements

    for (int q = 0; q < Q; q++) {
        char t;
        cin >> t;
        if (t == 'Q') {
            // query how many numbers in array between index left and right are divisible by num
            int left, right, num;
            cin >> left >> right >> num;
            int cnt = 0;
            for (int i = left; i <= right; i++) {
                if (array[i] % num == 0) {
                    cnt++;
                }
            }
            cout << cnt << '\n'; // don't use endl when there's a lot of output
        } else {
            // simple update on one element in array
            int index, newNum;
            cin >> index >> newNum;
            array[index] = newNum;
        }
    }
    return 0;
}

Answer 1

This can be solved by applying square root decomposition to make both range queries and point updates efficient. Let's split the array into approximately √N blocks each of size √N.

For each block, we can precompute the number of elements in the block divisible by each possible query value (the integers from 1 to 100000). We can find all factors of a number by iterating up to its square root, so we can add one element to the answer for a block in O(√(MAX_ARRAY_ELEMENT)). The precomputation for the initial array takes O(N√(MAX_ARRAY_ELEMENT)), which is acceptable since MAX_ARRAY_ELEMENT is at most 100000.

To update a value at an index, we remove the contribution of the current value at that index to its block, update the array, then add the contribution of the new element to the precomputed values for that block. This takes O(√(NEW_VALUE)).

To count the number of elements divisible by a particular number in a subarray, we can split the queried subarray into multiple segments. Some parts of the subarray might be directly covered by some blocks: there are at most √N of these and we can get the answer for each block in constant time as they are all precomputed. There are also at most two sections (at the beginning and end of the subarray) that only lie partially within a block; for these we need to use the brute force method to test all elements for divisibility, but there are only at most √N elements in each of these two blocks. The answer is the sum of the counts of all these sections. Thus, the time complexity for one such query is O(√N).

#include <iostream>
#include <vector>
#include <cmath>
#include <cstdlib>
constexpr int MAX_ELEM = 1e5;
int main() {
    int N, Q;
    std::cin >> N >> Q;
    int blockSize = std::sqrt(N), numBlocks = (N + blockSize - 1) / blockSize;
    std::vector<int> array;
    array.reserve(N);
    std::vector<std::vector<int>> blockAns(numBlocks, std::vector<int>(MAX_ELEM + 1));
    auto update = [&](int idx, int delta) {
        for (int i = 1, block = idx / blockSize; i * i <= array[idx]; ++i) {
            if (auto [q, r] = std::div(array[idx], i); r == 0) {
                blockAns[block][i] += delta;
                if (q != i) blockAns[block][q] += delta;
            }
        }
    };
    for (int i = 0, x; i < N; ++i) {
        std::cin >> x;
        array.push_back(x);
        update(i, 1);
    }
    for (int q = 0; q < Q; ++q) {
        char t;
        std::cin >> t;
        if (t == 'Q') {
            int left, right, num;
            std::cin >> left >> right >> num;
            int leftBlock = left / blockSize, rightBlock = right / blockSize, count = 0;
            if (leftBlock != rightBlock) {
                for (int i = left; i < blockSize * (leftBlock + 1); ++i) 
                    count += array[i] % num == 0;
                for (int b = leftBlock + 1; b < rightBlock; ++b)
                    count += blockAns[b][num];
                for (int i = rightBlock * blockSize; i <= right; ++i)
                    count += array[i] % num == 0;
            } else {
                for (int i = left; i <= right; ++i)
                    count += array[i] % num == 0;
            }
            std::cout << count << '\n';
        } else {
            int index;
            std::cin >> index;
            update(index, -1);
            std::cin >> array[index];
            update(index, 1);
        }
    }
}

Answer 2

Sorry for taking so long to write this up. I will skip the I/O.

from sortedcontainers import SortedList

def build_all_factors (n):
    # We'll simplify indexing by having a useless 0.
    all_factors = [[1] for _ in range(n+1)]
    for i in range(2, n+1):
        j = i
        while j <= n:
            all_factors[j].append(i)
            j += i
    return all_factors

def query_counts (n, values, queries):
    all_factors = build_all_factors(n)

    # Build up a data structure for what multiple is where.
    factor_indexes = [SortedList() for _ in range(n+1)]
    for i in range(len(values)):
        value = values[i]
        for factor in all_factors[value]:
            factor_indexes[factor].add(i)

    for query in queries:
        if query[0] == "Q":
            _, left, right, value = query
            index_left = factor_indexes[value].bisect_left(left)
            index_right = factor_indexes[value].bisect_right(right)
            print(index_right - index_left)
        else:
            _, i, new_value = query
            # Remove the existing value.
            value = values[i]
            for factor in all_factors[value]:
                factor_indexes[factor].remove(i)
            # Update and insert
            values[i] = new_value
            for factor in all_factors[new_value]:
                factor_indexes[factor].add(i)

query_counts(10, [6, 5, 4, 3, 2, 1], [
    ['Q', 0, 3, 1],
    ['Q', 1, 4, 2],
    ['U', 3, 8],
    ['Q', 1, 4, 2],
    ])

What does space and performance look like? Let's let n be the largest integer we can have, and m be the number of values that we have.

First, the full factorization table in all_factors takes space O(n log(n)) and takes the same time to build.

The initial factor_indexes with random integers requires space O(m log(n)). Each insert takes at most O(log(n)) time. (On average it's better than that, because most factors appear at only a small number of indexes. But that's fine for an upper bound.) But even with the most pathological possible input, using a highly composite number over and over again (in your case, 83160), which means 128 * m numbers to store. Even this isn't so bad.

A Q is never more than O(log(n)) to perform. And, on random operations, is better than that because there are probably not many factors.

An U takes longer. On average we have O(log(n)) removes and inserts, each of which takes at most O(log(n)). Leading to O(log(n)^2). The pathological case is again 128 deletes and 128 inserts. Which is again not so bad. Even your worst possible update should still be faster than a naive query was before.

Of course, I've pushed the real work down to SortedList. But as you mentioned in the comments, you can use C++ containers that allow you to take the same approach.