|
41 | 41 |
|
42 | 42 | <!-- 这里可写通用的实现逻辑 -->
|
43 | 43 |
|
44 |
| -树状数组。 |
| 44 | +**方法一:树状数组** |
45 | 45 |
|
46 | 46 | 树状数组,也称作“二叉索引树”(Binary Indexed Tree)或 Fenwick 树。 它可以高效地实现如下两个操作:
|
47 | 47 |
|
|
56 | 56 |
|
57 | 57 | 我们可以用树状数组,从左到右扫描前缀和数组,每遇到一个前缀和 s,就在树状数组中查询区间 `[preSum[j] - upper, preSum[j] - lower]` 内的整数的数量,随后将 s 更新至树状数组。
|
58 | 58 |
|
| 59 | +**方法二:线段树** |
| 60 | + |
| 61 | +Python3 代码 TLE,Java 代码 AC。 |
| 62 | + |
59 | 63 | <!-- tabs:start -->
|
60 | 64 |
|
61 | 65 | ### **Python3**
|
62 | 66 |
|
63 | 67 | <!-- 这里可写当前语言的特殊实现逻辑 -->
|
64 | 68 |
|
| 69 | +树状数组: |
| 70 | + |
65 | 71 | ```python
|
66 | 72 | class BinaryIndexedTree:
|
67 | 73 | def __init__(self, n):
|
@@ -106,10 +112,80 @@ class Solution:
|
106 | 112 | return ans
|
107 | 113 | ```
|
108 | 114 |
|
| 115 | +线段树: |
| 116 | + |
| 117 | +```python |
| 118 | +class Node: |
| 119 | + def __init__(self): |
| 120 | + self.l = 0 |
| 121 | + self.r = 0 |
| 122 | + self.v = 0 |
| 123 | + |
| 124 | +class SegmentTree: |
| 125 | + def __init__(self, n): |
| 126 | + self.tr = [Node() for _ in range(4 * n)] |
| 127 | + self.build(1, 1, n) |
| 128 | + |
| 129 | + def build(self, u, l, r): |
| 130 | + self.tr[u].l = l |
| 131 | + self.tr[u].r = r |
| 132 | + if l == r: |
| 133 | + return |
| 134 | + mid = (l + r) >> 1 |
| 135 | + self.build(u << 1, l, mid) |
| 136 | + self.build(u << 1 | 1, mid + 1, r) |
| 137 | + |
| 138 | + def modify(self, u, x, v): |
| 139 | + if self.tr[u].l == x and self.tr[u].r == x: |
| 140 | + self.tr[u].v += v |
| 141 | + return |
| 142 | + mid = (self.tr[u].l + self.tr[u].r) >> 1 |
| 143 | + if x <= mid: |
| 144 | + self.modify(u << 1, x, v) |
| 145 | + else: |
| 146 | + self.modify(u << 1 | 1, x, v) |
| 147 | + self.pushup(u) |
| 148 | + |
| 149 | + def pushup(self, u): |
| 150 | + self.tr[u].v = self.tr[u << 1].v + self.tr[u << 1 | 1].v |
| 151 | + |
| 152 | + def query(self, u, l, r): |
| 153 | + if self.tr[u].l >= l and self.tr[u].r <= r: |
| 154 | + return self.tr[u].v |
| 155 | + mid = (self.tr[u].l + self.tr[u].r) >> 1 |
| 156 | + v = 0 |
| 157 | + if l <= mid: |
| 158 | + v += self.query(u << 1, l, r) |
| 159 | + if r > mid: |
| 160 | + v += self.query(u << 1 | 1, l, r) |
| 161 | + return v |
| 162 | + |
| 163 | +class Solution: |
| 164 | + def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int: |
| 165 | + s = [0] |
| 166 | + for x in nums: |
| 167 | + s.append(s[-1] + x) |
| 168 | + alls = set() |
| 169 | + for v in s: |
| 170 | + alls.add(v) |
| 171 | + alls.add(v - lower) |
| 172 | + alls.add(v - upper) |
| 173 | + m = {v: i for i, v in enumerate(sorted(alls), 1)} |
| 174 | + tree = SegmentTree(len(m)) |
| 175 | + ans = 0 |
| 176 | + for v in s: |
| 177 | + l, r = m[v - upper], m[v - lower] |
| 178 | + ans += tree.query(1, l, r) |
| 179 | + tree.modify(1, m[v], 1) |
| 180 | + return ans |
| 181 | +``` |
| 182 | + |
109 | 183 | ### **Java**
|
110 | 184 |
|
111 | 185 | <!-- 这里可写当前语言的特殊实现逻辑 -->
|
112 | 186 |
|
| 187 | +树状数组: |
| 188 | + |
113 | 189 | ```java
|
114 | 190 | class Solution {
|
115 | 191 | public int countRangeSum(int[] nums, int lower, int upper) {
|
@@ -172,6 +248,102 @@ class BinaryIndexedTree {
|
172 | 248 | }
|
173 | 249 | ```
|
174 | 250 |
|
| 251 | +线段树: |
| 252 | + |
| 253 | +```java |
| 254 | +class Solution { |
| 255 | + public int countRangeSum(int[] nums, int lower, int upper) { |
| 256 | + int n = nums.length; |
| 257 | + long[] preSum = new long[n + 1]; |
| 258 | + for (int i = 0; i < n; ++i) { |
| 259 | + preSum[i + 1] = preSum[i] + nums[i]; |
| 260 | + } |
| 261 | + TreeSet<Long> ts = new TreeSet<>(); |
| 262 | + for (long s : preSum) { |
| 263 | + ts.add(s); |
| 264 | + ts.add(s - upper); |
| 265 | + ts.add(s - lower); |
| 266 | + } |
| 267 | + Map<Long, Integer> m = new HashMap<>(); |
| 268 | + int idx = 1; |
| 269 | + for (long s : ts) { |
| 270 | + m.put(s, idx++); |
| 271 | + } |
| 272 | + int ans = 0; |
| 273 | + SegmentTree tree = new SegmentTree(m.size()); |
| 274 | + for (long s : preSum) { |
| 275 | + int l = m.get(s - upper); |
| 276 | + int r = m.get(s - lower); |
| 277 | + ans += tree.query(1, l, r); |
| 278 | + tree.modify(1, m.get(s), 1); |
| 279 | + } |
| 280 | + return ans; |
| 281 | + } |
| 282 | +} |
| 283 | + |
| 284 | +class Node { |
| 285 | + int l; |
| 286 | + int r; |
| 287 | + int v; |
| 288 | +} |
| 289 | + |
| 290 | +class SegmentTree { |
| 291 | + private Node[] tr; |
| 292 | + |
| 293 | + public SegmentTree(int n) { |
| 294 | + tr = new Node[4 * n]; |
| 295 | + for (int i = 0; i < tr.length; ++i) { |
| 296 | + tr[i] = new Node(); |
| 297 | + } |
| 298 | + build(1, 1, n); |
| 299 | + } |
| 300 | + |
| 301 | + public void build(int u, int l, int r) { |
| 302 | + tr[u].l = l; |
| 303 | + tr[u].r = r; |
| 304 | + if (l == r) { |
| 305 | + return; |
| 306 | + } |
| 307 | + int mid = (l + r) >> 1; |
| 308 | + build(u << 1, l, mid); |
| 309 | + build(u << 1 | 1, mid + 1, r); |
| 310 | + } |
| 311 | + |
| 312 | + public void modify(int u, int x, int v) { |
| 313 | + if (tr[u].l == x && tr[u].r == x) { |
| 314 | + tr[u].v += v; |
| 315 | + return; |
| 316 | + } |
| 317 | + int mid = (tr[u].l + tr[u].r) >> 1; |
| 318 | + if (x <= mid) { |
| 319 | + modify(u << 1, x, v); |
| 320 | + } else { |
| 321 | + modify(u << 1 | 1, x, v); |
| 322 | + } |
| 323 | + pushup(u); |
| 324 | + } |
| 325 | + |
| 326 | + public void pushup(int u) { |
| 327 | + tr[u].v = tr[u << 1].v + tr[u << 1 | 1].v; |
| 328 | + } |
| 329 | + |
| 330 | + public int query(int u, int l, int r) { |
| 331 | + if (tr[u].l >= l && tr[u].r <= r) { |
| 332 | + return tr[u].v; |
| 333 | + } |
| 334 | + int mid = (tr[u].l + tr[u].r) >> 1; |
| 335 | + int v = 0; |
| 336 | + if (l <= mid) { |
| 337 | + v += query(u << 1, l, r); |
| 338 | + } |
| 339 | + if (r > mid) { |
| 340 | + v += query(u << 1 | 1, l, r); |
| 341 | + } |
| 342 | + return v; |
| 343 | + } |
| 344 | +} |
| 345 | +``` |
| 346 | + |
175 | 347 | ### **C++**
|
176 | 348 |
|
177 | 349 | ```cpp
|
|
0 commit comments