原题链接在这里:https://leetcode.com/problems/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
.
题解:
题目的意思是说给了一个int array, 计算有多少subarray的sum在[lower, upper]区间内. 给的例子是index.
建立BST,每个TreeNode的val是prefix sum. 为了避免重复的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大的数目,再从总数里面减去,就是所求
Time Complexity: O(nlogn). Space: O(n).
AC Java:
1 public class Solution { 2 public int countRangeSum(int[] nums, int lower, int upper) { 3 if(nums == null || nums.length == 0){ 4 return 0; 5 } 6 int res = 0; 7 long [] sum = new long[nums.length+1]; 8 for(int i = 1; i<sum.length; i++){ 9 sum[i] = sum[i-1] + nums[i-1]; 10 } 11 12 TreeNode root = new TreeNode(sum[0]); 13 for(int i = 1; i<sum.length; i++){ 14 res += rangeSize(root, sum[i]-upper, sum[i]-lower); 15 insert(root, sum[i]); 16 } 17 return res; 18 } 19 20 private TreeNode insert(TreeNode root, long val){ 21 if(root == null){ 22 return new TreeNode(val); 23 } 24 if(root.val == val){ 25 root.count++; 26 }else if(root.val > val){ 27 root.leftSize++; 28 root.left = insert(root.left, val); 29 }else if(root.val < val){ 30 root.rightSize++; 31 root.right = insert(root.right, val); 32 } 33 return root; 34 } 35 36 private int countSmaller(TreeNode root, long val){ 37 if(root == null){ 38 return 0; 39 } 40 if(root.val == val){ 41 return root.leftSize; 42 }else if(root.val > val){ 43 return countSmaller(root.left, val); 44 }else{ 45 return root.leftSize + root.count + countSmaller(root.right, val); 46 } 47 } 48 49 private int countLarget(TreeNode root, long val){ 50 if(root == null){ 51 return 0; 52 } 53 if(root.val == val){ 54 return root.rightSize; 55 }else if(root.val > val){ 56 return countLarget(root.left, val) + root.count + root.rightSize; 57 }else{ 58 return countLarget(root.right, val); 59 } 60 } 61 62 private int rangeSize(TreeNode root, long lower, long upper){ 63 int total = root.leftSize + root.count + root.rightSize; 64 int smaller = countSmaller(root, lower); 65 int larger = countLarget(root, upper); 66 return total - smaller - larger; 67 } 68 } 69 70 class TreeNode{ 71 long val; 72 int count; 73 int leftSize; 74 int rightSize; 75 TreeNode left; 76 TreeNode right; 77 public TreeNode(long val){ 78 this.val = val; 79 this.count = 1; 80 this.leftSize = 0; 81 this.rightSize = 0; 82 } 83 }