How to find sum of maximum K elements in range in array

Question

Recently, I came up to the following problem:

Given array $A$ of size $n$ and integer $k$, We should answer $Q$ questions of the type: for given range $[l, r]$ we should find sum of the $k$ maximum elements over all $A_i$ such that $l \leq i \leq r$.

Trivial solution would be to sort all values in the range and iterate over the k biggest elements, however this has complexity $O(QN\log N)$. After some searching on the internet, I found out that this can be solved by segment trees, however I couldn't think of proper use of segment trees for this problem.

Is it possible to solve this problem in complexity better than $O(QN\log N)$

Answer 1

This answer refers to the version of the question in which the interval $[l,r]$ refers to the values of the elements rather than their indices.

Without preprocessing: You can extract all elements in the range $[l,r]$ in linear time. You can then use the linear time selection algorithm to find the $k$'th largest element, and then you can easily find the sum of the $k$ maximum elements. Overall, you can answer each query in $O(n)$.

With preprocessing: Sort the array, and compute running sums. Given $l,r$, locate the subset of the array corresponding to this range using binary search, and then use the running sums to compute the sum of the $k$ largest elements in $O(1)$. Overall, this uses $O(n\log n)$ preprocessing and $O(\log n)$ time per query.

Answer 2

This answer refers to the version of the question in which the range $[l,r]$ refers to indices of the array, and in which we have $Q$ queries. The question asked whether we could beat $O(Qn\log n)$.

Without preprocessing: Using the linear time selection algorithm, you can answer each query in $O(n)$, for a total of $O(Qn)$, which is better than $O(Qn\log n)$.

With preprocessing: Assume for simplicity that $n$ is a power of 2. We consider $\log n$ layers, the first consisting of $n$ intervals of length 1, the second of $n/2$ intervals of length 2, and so on: the $k$th consists of $n/2^k$ intervals of length $k$. We sort all of the intervals in $O(n\log n)$ by "merging up the tree". At query time, we partition the input interval into $O(\log n)$ stored intervals, and them return the answer in $O(k\log n)$ by merging the $k$th largest elements from each.

We can improve the query time as follows. Along with the layered intervals, we also compute their running sums. Let us consider the $O(\log n)$ layered intervals $I_1,\ldots,I_m$ whose union is the given interval $[l,r]$. The $k$ largest elements of $[l,r]$ are composed of the $k_1,\ldots,k_m$ largest elements in $I_1,\ldots,I_m$. If we could determine $k_1,\ldots,k_m$, then we could compute the answer in $O(m) = O(\log n)$ using the running sums.

We can determine $k_1,\ldots,k_m$ as follows. For each element $x \in I_1$, we can count how many elements are larger than $x$ in time $O(m\log k) = O(\log n \log k)$ using binary search. By doing binary search on the top $k$ elements of $I_1$, we can determine $k_1$ in time $O(\log n \log^2 k)$. Continuing in this way, we can determine $k_1,\ldots,k_m$ in time $O(\log^2 n \log^2 k)$.

This solution has preprocessing time $O(n\log n)$ and query time $O(\log^2 n \log^2 k)$.

We can obtain other tradeoffs by varying the fan-in of the tree. Suppose that the fan-in is $n^{1/d}$. For each internal node in the tree, we sort each interval-of-intervals within its children. At level $i$ we have $n^{i/d}$ vertices, each has $n^{2/d}$ intervals-of-intervals, each of length at most $n^{1-i/d}$, so sorting all intervals-of-intervals at level $i$ takes time $n^{1+2/d}$, for a total preprocessing time of $O(dn^{1+2/d})$. In query time, we can break $[l,r]$ into a union of $O(d)$ intervals, for a total query time of $O(d^2\log^2 k)$.

Stated differently, for each $\epsilon > 0$ this gives an algorithm with preprocessing time $O(\epsilon^{-1} n^{1+\epsilon})$ and query time $O(\epsilon^{-2} \log^2 k)$.