aboutsummaryrefslogtreecommitdiff
path: root/range_sum_sorted_subarr.py
blob: 81e28cec528d22af08abe66e95dcd3630d2033a6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
# 1508. Range Sum of Sorted Subarray Sums

"""
you are given the array 'nums' consisting of 'n' positive integers. you
computed the sum of all non empty continuous subarrays from the array and
then sorted them in non decreasing order, creating a new array of 'n  (n +
1) / 2' numbers. return the sum of the numbers from index 'left' to index
'right' inclusive, in the new array. since the answer may be very large,
return it modulo 1e9+7
"""


class Solution(object):
    def rangeSum(self, nums, n, left, right):
        """
        :type nums: List[int]
        :type n: int
        :type left: int
        :type right: int
        :rtype: int
        """
        b, c = [0] * (n + 1), [0] * (n + 1)
        for i in range(n):
            b[i + 1] = b[i] + nums[i]
            c[i + 1] = c[i] + b[i + 1]

        def cnt_sum_und(score):
            res, i = 0, 0
            for j in range(n + 1):
                while b[j] - b[i] > score:
                    i += 1
                res += j - i
            return res

        def kth_score(k):
            l, r = 0, b[n]
            while l < r:
                m = (l + r) / 2
                if cnt_sum_und(m) < k:
                    l = m + 1
                else:
                    r = m
            return l

        def sum_k_sums(k):
            score = kth_score(k)
            res, i = 0, 0
            for j in range(n + 1):
                while b[j] - b[i] > score:
                    i += 1
                res += b[j] * (j - i + 1) - (c[j] - (c[i - 1] if i else 0))
            return res - (cnt_sum_und(score) - k) * score

        return sum_k_sums(right) - sum_k_sums(left - 1)


if __name__ == "__main__":
    obj = Solution()
    print(obj.rangeSum(nums=[1, 2, 3, 4], n=4, left=1, right=5))
    print(obj.rangeSum(nums=[1, 2, 3, 4], n=4, left=3, right=4))
    print(obj.rangeSum(nums=[1, 2, 3, 4], n=4, left=1, right=10))