1. 程式人生 > 實用技巧 >327. Count of Range Sum

327. Count of Range Sum

問題:

給定一個數組,求連續元素之和在給定範圍[lower, upper]之間的,連續idx為(i~j)元素組個數。

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

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

Constraints:
0 <= nums.length <= 10^4

  

解法:FenwickTree

解法類似315. Count of Smaller Numbers After Self

  • 只是需要多求一步presum->sum,作為處理物件。
  • 對字首和陣列,進行sort->sortsum,來確定在FenwickTree中的idx。
  • FenwickTree記錄對應相同sum值的個數(出現頻率)。

然後對原順序的sum,逐一讀入,

每個sum[i]中,在FenwickTree中,尋找sum[i]-upper ~ sum[i]-lower 範圍內共有多少個,res追加。

然後把當前sum[i]找到idx,插入FenwickTree中,累計count+1。

程式碼參考:

 1 class FenwickTree {
 2 public:
 3     FenwickTree(int n):tree(n+1, 0) {}
 4     void update(int i, int delta) {
 5         while(i<tree.size()) {
 6             tree[i]+=delta;
 7             i+=lowbit(i);
 8         }
 9     }
10     int getPresum(int i) {
11         int sum=0;
12         while
(i>0){ 13 sum+=tree[i]; 14 i-=lowbit(i); 15 } 16 return sum; 17 } 18 private: 19 vector<int> tree; 20 int lowbit(int x) { 21 return x&(-x); 22 } 23 }; 24 25 class Solution { 26 public: 27 int countRangeSum(vector<int>& nums, int lower, int upper) { 28 int cursum=0; 29 int res=0; 30 vector<long long> sum(nums.size()+1,0); 31 for(int i=0; i<nums.size(); i++){ 32 sum[i+1]=sum[i]+nums[i]; 33 } 34 vector<long long> sortsum(sum); 35 sort(sortsum.begin(), sortsum.end()); 36 37 FenwickTree ftree(sum.size()); 38 for(int i=0; i<sum.size(); i++){ 39 int idx_low = distance(sortsum.begin(), lower_bound(sortsum.begin(),sortsum.end(),sum[i]-upper)); 40 int idx_upper = distance(sortsum.begin(), upper_bound(sortsum.begin(),sortsum.end(),sum[i]-lower)); 41 res+=(ftree.getPresum(idx_upper)-ftree.getPresum(idx_low)); 42 43 int idx = distance(sortsum.begin(), lower_bound(sortsum.begin(),sortsum.end(),sum[i])); 44 ftree.update(idx+1, 1); 45 } 46 return res; 47 } 48 };