Can Merge Sort Tree find sum of K-th largest elements, in range [l, r] if the elements are not unique?

Question

See below code:

vector<long long> temp;
for (int i = l; i <= r; ++i)
{
    temp.push_back(holes[i]);
}
sort(temp.rbegin(), temp.rend());
long long needed_H2 = 0;
for (int i = K; i < temp.size(); ++i)
{
    needed_H2 += temp[i];
}

Basically, it copies some range [l, r] of original array, sorts it in non-ascending order and then finds sum of this range omitting first K greatest element.

$N$ is size of array on which we will make range queries.

I need to achieve the same in like $O(log^2 N)$ time complexity (if possible), and I wonder about Merge Sort Tree (MST) specifically. If the values would be unique - then I see solution in $O(log^2 N * log M)$ where $M$ is maximum value we will ever potentially check in binary search. The $log^2 N$ factor is due to traversing Merge Sort Tree, and at every node which is in range we look for, we will also do binary search (lower bound) to find first value greater than $X$ given by binary search. So, we perform binary search on the MST asking "how many elements are greater than X" and if we find X such that returned count is equal to K (it must happen if elements are unique and we have more elements than K). To find sum of these K-th largest elements, on MST we can also store prefix sums of sorted ranges which every MST node already stores.

But what if the element are not unique? Can we also solve it?

Lets say count of elements greater than given val is C.

I tried finding minimum value such that the $C >= K$, and then do some adjustmens using minimum value (because I think that if we would exclude all minimum elements from range, so we would search for a minimally greater val such that the C will change,then we would find $C < K$, because the C we found before was the minimal greater than K). At the end, we would substract from sum found on range the (minimum value * ($C - K$))

I wrote code, but unfortunately, it doesnt work as expected. I spent much time debugging, so I think there's something wrong this above idea.

So, can we find sum of first K largest elements in range [l, r] using Merge Sort Tree?

score 1 · Accepted Answer · answered Jan 11 '25 at 22:26

Suppose you have your merge sort tree built. I believe you can, in $\log^2(n)$ time, for given $x$ find the number of elements in range $[l,r]$ that are $\geq x$. Seeing this as a function of $x$, it is monotonic. I.e. for $y\geq x$ there cannot be less elements $\geq y$ than $\geq x$. So at this point I think you just want to look for the boundary $x'$ such that $|\{y \in [l,r] \wedge y \geq x'\}| \geq k$ and $|\{y\in [l,r] \wedge y \geq x' + 1\}|$ < k. Now, you know that you need to sum everything after (and including) $x'+1$ in the range, which you can do with your prefix sum. And then you need to find the remainder to k.

You can still achieve the time complexity $\log^2(n)\cdot \log M$ or even $\log^3(n)$ because you can binary search only over the at most $n$ numbers in the array.

Does this help?

Can Merge Sort Tree find sum of K-th largest elements, in range [l, r] if the elements are not unique?

1 Answers1