Leetcode: 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.

參考：https://leetcode.com/discuss/79083/share-my-solution

First of all, let's look at the naive solution. Preprocess to calculate the prefix sums S[i] = S(0, i), then S(i, j) = S[j] - S[i]. With these prefix sums, it is trivial to see that with O(n^2)time we can find all S(i, j) in the range [lower, upper]

Java - Naive Solution(這個做法為了讓所有區間都能表示成一個區間減另一個區間，size額外增加了1，sums[i]定義為前i個元素之和，這樣連只包含第一個元素的區間的和都可以表示為sums[1]-sums[0],這樣寫不用再分類討論，挺好的)

public int countRangeSum(int[] nums, int lower, int upper) { int n = nums.length; long[] sums = new long[n + 1]; for (int i = 0; i < n; ++i) sums[i + 1] = sums[i] + nums[i]; int ans = 0; for (int i = 0; i < n; ++i) for (int j = i + 1; j <= n; ++j) if (sums[j] - sums[i] >= lower && sums[j] - sums[i] <= upper) ans++; return ans; } 

However the naive solution is set to TLE intentionally

Now let's do better than this.

Recall count smaller number after self where we encountered the problem

• count[i] = count of nums[j] - nums[i] < 0 with j > i

Here, after we did the preprocess, we need to solve the problem

• count[i] = count of a <= S[j] - S[i] <= b with j > i
• ans = sum(count[:])

Therefore the two problems are almost the same. We can use the same technique used in that problem to solve this problem. One solution is merge sort based; another one is Balanced BST based. The time complexity are both O(n log n).

The merge sort based solution counts the answer while doing the merge. During the merge stage, we have already sorted the left half [start, mid) and right half [mid, end). We then iterate through the left half with index i. For each i, we need to find two indices k and j in the right half where

• j is the first index satisfy sums[j] - sums[i] > upper and
• k is the first index satisfy sums[k] - sums[i] >= lower.

Then the number of sums in [lower, upper] is j-k. We also use another index t to copy the elements satisfy sums[t] < sums[i] to a cache in order to complete the merge sort.

Despite the nested loops, the time complexity of the "merge & count" stage is still linear. Because the indices k, j, t will only increase but not decrease, each of them will only traversal the right half once at most. The total time complexity of this divide and conquer solution is then O(n log n).

One other concern is that the sums may overflow integer. So we use long instead.

 

方法一：mergesort, O(NlogN) running time 最快但是不喜歡這種寫法，不理解11行

public class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] sums = new long[n + 1];
        for (int i = 0; i < n; ++i)
            sums[i + 1] = sums[i] + nums[i];
        return countWhileMergeSort(sums, 0, n + 1, lower, upper);
    }
    
    private int countWhileMergeSort(long[] sums, int start, int end, int lower, int upper) {
        if (end - start <= 1) return 0;
        int mid = (start + end) / 2;
        int count = countWhileMergeSort(sums, start, mid, lower, upper) 
                  + countWhileMergeSort(sums, mid, end, lower, upper);
        int j = mid, k = mid, t = mid, r = 0;
        long[] cache = new long[end - start];
        for (int i = start; i < mid; ++i, ++r) {
            while (k < end && sums[k] - sums[i] < lower) k++;
            while (j < end && sums[j] - sums[i] <= upper) j++;
            while (t < end && sums[t] < sums[i]) cache[r++] = sums[t++]; //start merging
            cache[r] = sums[i];
            count += j - k;
        }
        System.arraycopy(cache, 0, sums, start, r);
        return count;
    }
}

 默認方法：construct BST (好理解很多) ， Time: O(NlogN)

這個做法是建立BST，把prefix sum作為TreeNode.val存進去，為了避免重復的TreeNode.val處理麻煩，設置一個count記錄多少個重復TreeNode.val, 維護leftSize, 記錄比該節點value小的節點個數，rightSize同理

由於RangeSum S(i,j)在[lower,upper]之間的條件是lower<=sums[j+1]-sums[i]<=upper, 所以我們每次insert一個新的PrefixSum sums[k]進這個BST之前，先尋找一下（rangeSize）該BST內已經有多少個PrefixSum（叫它sums[t]吧）滿足lower<=sums[k]-sums[t]<=upper, 即尋找有多少個sums[t]滿足： 

sums[k]-upper<=sums[t]<=sums[k]-lower

BST提供了countSmaller和countLarger的功能，計算比sums[k]-upper小的RangeSum數目和比sums[k]-lower大的數目，再從總數里面減去，就是所求

public class Solution {
    private class TreeNode {
        long val = 0;
        int count = 1;
        int leftSize = 0;
        int rightSize = 0;
        TreeNode left = null;
        TreeNode right = null;
        public TreeNode(long v) {
            this.val = v;
            this.count = 1;
            this.leftSize = 0;
            this.rightSize = 0;
        }
    }

    private TreeNode insert(TreeNode root, long val) {
        if(root == null) {
            return new TreeNode(val);
        } else if(root.val == val) {
            root.count++;
        } else if(val < root.val) {
            root.leftSize++;
            root.left = insert(root.left, val);
        } else if(val > root.val) {
            root.rightSize++;
            root.right = insert(root.right, val);
        }
        return root;
    }

    private int countSmaller(TreeNode root, long val) {
        if(root == null) {
            return 0;
        } else if(root.val == val) {
            return root.leftSize;
        } else if(root.val > val) {
            return countSmaller(root.left, val);
        } else {
            return root.leftSize + root.count + countSmaller(root.right, val);
        }
    }

    private int countLarger(TreeNode root, long val) {
        if(root == null) {
            return 0;
        } else if(root.val == val) {
            return root.rightSize;
        } else if(root.val < val) {
            return countLarger(root.right, val);
        } else {
            return countLarger(root.left, val) + root.count + root.rightSize;
        }
    }

    private int rangeSize(TreeNode root, long lower, long upper) {
        int total = root.count + root.leftSize + root.rightSize;
        int smaller = countSmaller(root, lower);    // Exclude everything smaller than lower
        int larger = countLarger(root, upper);      // Exclude everything larger than upper
        return total - smaller - larger;
    }

    public int countRangeSum(int[] nums, int lower, int upper) {
        if(nums.length == 0) {
            return 0;
        }
        long[] sums = new long[nums.length + 1];
        for(int i = 0; i < nums.length; i++) {
            sums[i + 1] = sums[i] + nums[i];
        }
        TreeNode root = new TreeNode(sums[0]);
        int output = 0;
        for(int i = 1; i < sums.length; i++) {
            output += rangeSize(root, sums[i] - upper, sums[i] - lower);
            insert(root, sums[i]);
        }
        return output;
    }
}