2017-09-14 LeetCode_327 Count of Range Sum

327. Count of Range Sum

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.
Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.

Example:
Given nums = [-2, 5, -1], lower = -2, upper = 2,
Return 3.
The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

Credits:
Special thanks to @dietpepsi for adding this problem and creating all test cases.

Solution:

class Solution {

2

public:

3

    void sort1(vector<long long>& sums, vector<long long>& nums, long long start, long long end) {

4

        if (start >= end) return;

5

        long long m1 = sums[start], m2 = nums[start], i = start, j = end;

6

        while (i < j) {

7

            while (i < j && sums[j]-nums[j] >= m1-m2) j--;

8

            sums[i] = sums[j]; nums[i] = nums[j];

9

            while (i < j && sums[i]-nums[i] <= m1-m2) i++;

10

            sums[j] = sums[i]; nums[j] = nums[i];

11

        }

12

        sums[i] = m1; nums[i] = m2;

13

        sort1(sums, nums, start, i-1);

14

        sort1(sums, nums, i+1, end);

15

    }

16

    void sort2(vector<long long>& sums, vector<long long>& nums, long long start, long long end) {

17

        if (start >= end) return;

18

        long long m1 = sums[start], m2 = nums[start], i = start, j = end;

19

        while (i < j) {

20

            while (i < j && sums[j] >= m1) j--;

21

            sums[i] = sums[j]; nums[i] = nums[j];

22

            while (i < j && sums[i] <= m1) i++;

23

            sums[j] = sums[i]; nums[j] = nums[i];

24

        }

25

        sums[i] = m1; nums[i] = m2;

26

        sort2(sums, nums, start, i-1);

27

        sort2(sums, nums, i+1, end);

28

    }

29

    long long divides(vector<long long>& sums, vector<long long>& nums, long long start, long long end, int lower, int upper) {

30

        if (start == end) return (nums[start] <= upper && nums[start] >= lower);

31

        long long count = divides(sums, nums, start, (start+end)/2, lower, upper)+divides(sums, nums, (start+end)/2+1, end, lower, upper);

32

        sort1(sums, nums, start, (start+end)/2);

33

        sort2(sums, nums, (start+end)/2+1, end);

34

35

        for (int left = start, right0 = (start+end)/2+1, right1 = (start+end)/2+1; left <= (start+end)/2; left++) {

36

            while (right0 <= end && sums[right0]-sums[left]+nums[left] < lower) right0++;

37

            while (right1 <= end && sums[right1]-sums[left]+nums[left] <= upper) right1++;

38

            count += right1-right0;

39

        }

40

        return count;

41

    }

42

    int countRangeSum(vector<int>& nums, int lower, int upper) {

43

        if (nums.size() == 0) return 0;

44

        vector<long long> sums, mynums; sums.push_back(nums[0]); mynums.push_back(nums[0]);

45

        for (int i = 1; i < nums.size(); i++) {

46

            mynums.push_back(nums[i]);

47

            sums.push_back(sums[i-1]+nums[i]);

48

        }

49

        return divides(sums, mynums, 0, nums.size()-1, lower, upper);

50

    }

51

};