Search code examples
c++algorithmbinary-searchfenwick-treebinary-indexed-tree

Fenwick tree(BIT). Find the smallest index with given cumulative frequency in O(logN)


let's say I have a BIT(Fenwick Tree) with non-negative values and I want to find the smallest index in it for given cumulative frequency in O(logN).

Now, I can do it O(log^2(N)) like this.

int l = 0, r = n;
while(l < r) {
    int midd = l + (r-l)/2;
    if (get_sum(BIT, midd+1) < given_sum)
        l = midd+1;
    else
        r = midd;
}
return midd + 1;

I know that we can find any of index or the greatest index in O(logN) as described here, so expect to find the lowest with the same time complexity as well. The way the tree is implemented is a common one.

vector<int> BIT(n+1);
void update(vector<int> &BIT, int idx, int delta){
    for(int i = idx; i < BIT.size(); i +=(i&-i))
        BIT[i] += delta;
}
int get_sum(vector<int>& BIT, int idx){
    int sum = 0;
    for(int i = idx; i > 0; i-=(i&-i))
        sum += BIT[i];
    return sum;
}

Hope for your help:)


Solution

  • Here is my implementation of a lower_bound-like function for a Fenwick tree with 0-based indexing:

    std::size_t lower_bound(int value) const
    {
        std::size_t index = 0;
        for (auto mask = msb_size_mask(); mask != 0; mask >>= 1)
        {
            const auto k = mask + index - 1;
            if (k < data.size() && data[k] < value)
            {
                value -= data[k];
                index += mask;
            }
        }
        return index;
    }
    

    data is the underlying std::vector<int>. The helper function msb_size_mask() returns the minimum size of a Fenwick tree such that the underlying binary tree is perfect, i.e. it returns 2^k if data.size() is in the range (2^{k-1}, 2^k]. In C++20, this is exactly what std::bit_ceil() does.

    Here is the update function:

    void add(std::size_t index, int value)
    {
        for (; index < data.size(); index |= index + 1)
            data[index] += value;
    }
    

    The complete code can be found here.