地址 https://leetcode-cn.com/problems/count-of-range-sum/
给定一个整数数组 nums,返回区间和在 [lower, upper] 之间的个数, 包含 lower 和 upper。 区间和 S(i, j) 表示在 nums 中, 位置从 i 到 j 的元素之和,包含 i 和 j (i ≤ j)。 说明: 最直观的算法复杂度是 O(n2) ,请在此基础上优化你的算法。 示例: 输入: nums = [-2,5,-1], lower = -2, upper = 2, 输出: 3 解释: 3个区间分别是: [0,0], [2,2], [0,2],它们表示的和分别为: -2, -1, 2。
算法1
典型的连续空间求和模板
Y总得模板更好 hh
1 线段树
C++ 代码
struct SegNode { int lo, hi, add; SegNode* lchild, *rchild; SegNode(int left, int right): lo(left), hi(right), add(0), lchild(nullptr), rchild(nullptr) {} }; class Solution { public: SegNode* build(int left, int right) { SegNode* node = new SegNode(left, right); if (left == right) { return node; } int mid = (left + right) / 2; node->lchild = build(left, mid); node->rchild = build(mid + 1, right); return node; } void insert(SegNode* root, int val) { root->add++; if (root->lo == root->hi) { return; } int mid = (root->lo + root->hi) / 2; if (val <= mid) { insert(root->lchild, val); } else { insert(root->rchild, val); } } int count(SegNode* root, int left, int right) const { if (left > root->hi || right < root->lo) { return 0; } if (left <= root->lo && root->hi <= right) { return root->add; } return count(root->lchild, left, right) + count(root->rchild, left, right); } int countRangeSum(vector<int>& nums, int lower, int upper) { long long sum = 0; vector<long long> preSum = {0}; for (int v: nums) { sum += v; preSum.push_back(sum); } set<long long> allNumbers; for (long long x: preSum) { allNumbers.insert(x); allNumbers.insert(x - lower); allNumbers.insert(x - upper); } // 利用哈希表进行离散化 unordered_map<long long, int> values; int idx = 0; for (long long x: allNumbers) { values[x] = idx; idx++; } SegNode* root = build(0, values.size() - 1); int ret = 0; for (long long x: preSum) { int left = values[x - upper], right = values[x - lower]; ret += count(root, left, right); insert(root, values[x]); } return ret; } };
算法2
2树状数组
C++ 代码
class BIT { private: vector<int> tree; int n; public: BIT(int _n): n(_n), tree(_n + 1) {} static constexpr int lowbit(int x) { return x & (-x); } void update(int x, int d) { while (x <= n) { tree[x] += d; x += lowbit(x); } } int query(int x) const { int ans = 0; while (x) { ans += tree[x]; x -= lowbit(x); } return ans; } }; class Solution { public: int countRangeSum(vector<int>& nums, int lower, int upper) { long long sum = 0; vector<long long> preSum = {0}; for (int v: nums) { sum += v; preSum.push_back(sum); } set<long long> allNumbers; for (long long x: preSum) { allNumbers.insert(x); allNumbers.insert(x - lower); allNumbers.insert(x - upper); } // 利用哈希表进行离散化 unordered_map<long long, int> values; int idx = 0; for (long long x: allNumbers) { values[x] = idx; idx++; } int ret = 0; BIT bit(values.size()); for (int i = 0; i < preSum.size(); i++) { int left = values[preSum[i] - upper], right = values[preSum[i] - lower]; ret += bit.query(right + 1) - bit.query(left); bit.update(values[preSum[i]] + 1, 1); } return ret; } };