Range Sum of Sorted Subarray Sums

lcpy gen -n 1508   # by problem number
lcpy gen -s range_sum_of_sorted_subarray_sums   # by problem name

Input: nums = [1,2,3,4], n = 4, left = 1, right = 5
Output: 13
Explanation: All subarray sums are 1, 3, 6, 10, 2, 5, 9, 3, 7, 4. After sorting them in non-decreasing order we have the new array [1, 2, 3, 3, 4, 5, 6, 7, 9, 10]. The sum of the numbers from index le = 1 to ri = 5 is 1 + 2 + 3 + 3 + 4 = 13.

Input: nums = [1,2,3,4], n = 4, left = 3, right = 4
Output: 6
Explanation: The given array is the same as example 1. We have the new array [1, 2, 3, 3, 4, 5, 6, 7, 9, 10]. The sum of the numbers from index le = 3 to ri = 4 is 3 + 3 = 6.

Input: nums = [1,2,3,4], n = 4, left = 1, right = 10
Output: 50

class Solution:
    # Time: O(n * log(sum(nums))) ~ O(n log n)
    # Space: O(1)
    def range_sum(self, nums: list[int], n: int, left: int, right: int) -> int:
        mod = 1_000_000_007

        def count_and_sum(target: int) -> tuple[int, int]:
            # Number of subarrays with sum <= target, plus the sum of those sums.
            count = 0
            total = 0
            window_sum = 0
            window_total = 0
            start = 0
            for end, value in enumerate(nums):
                window_sum += value
                window_total += value * (end - start + 1)
                while window_sum > target:
                    window_total -= window_sum
                    window_sum -= nums[start]
                    start += 1
                count += end - start + 1
                total += window_total
            return count, total

        def prefix_sum_of_sums(k: int) -> int:
            # Sum of the k smallest subarray sums.
            lo, hi = min(nums), sum(nums)
            while lo < hi:
                mid = (lo + hi) // 2
                count, _ = count_and_sum(mid)
                if count < k:
                    lo = mid + 1
                else:
                    hi = mid
            count, total = count_and_sum(lo)
            return total - lo * (count - k)

        return (prefix_sum_of_sums(right) - prefix_sum_of_sums(left - 1)) % mod