|
45 | 45 |
|
46 | 46 | <!-- 这里可写通用的实现逻辑 -->
|
47 | 47 |
|
| 48 | +本题可以用树状数组解决。 |
| 49 | + |
| 50 | +树状数组,也称作“二叉索引树”(Binary Indexed Tree)或 Fenwick 树。 它可以高效地实现如下两个操作: |
| 51 | + |
| 52 | +1. **单点更新** `update(x, delta)`: 把序列 x 位置的数加上一个值 delta; |
| 53 | +1. **前缀和查询** `query(x)`:查询序列 `[1,...x]` 区间的区间和,即位置 x 的前缀和。 |
| 54 | + |
| 55 | +这两个操作的时间复杂度均为 `O(log n)`。 |
| 56 | + |
| 57 | +对于本题,我们先用 pos 记录每个数在 nums2 中的位置,然后依次对 nums1 中的每个元素进行处理。 |
| 58 | + |
| 59 | +考虑**以当前数字作为三元组中间数字**的好三元组的数目。第一个数字需要是之前已经遍历过的,并且在 nums2 中的位置比当前数字更靠前的;第三个数字需要是当前还没有遍历过的,并且在 nums2 中的位置比当前数字更靠后的。 |
| 60 | + |
| 61 | +以 `nums1 = [4,0,1,3,2], nums2 = [4,1,0,2,3]`为例,考虑我们的遍历过程: |
| 62 | + |
| 63 | +1. 首先处理 4,此时 nums2 中出现情况为 `[4,X,X,X,X]`,4 之前有值的个数是 0,4 之后没有值的个数有 4 个。因此以 4 为中间数字能形成 0 个好三元组。 |
| 64 | +1. 接下来是 0,此时 nums2 中出现情况为 `[4,X,0,X,X]`,0 之前有值的个数是 1,0 之后没有值的个数有 2 个。因此以 0 为中间数字能形成 2 个好三元组。 |
| 65 | +1. 接下来是 1,此时 nums2 中出现情况为 `[4,1,0,X,X]`,1 之前有值的个数是 1,0 之后没有值的个数有 2 个。因此以 1 为中间数字能形成 2 个好三元组。 |
| 66 | +1. ... |
| 67 | +1. 最后是 2,此时 nums2 中出现情况为 `[4,1,0,2,3]`,2 之前有值的个数是 4,2 之后没有值的个数是 0。因此以 2 为中间数字能形成 0 个好三元组。 |
| 68 | + |
| 69 | +我们可以用树状数组来更新 nums2 中各个位置数字的出现情况,快速算出每个数字左侧 1 的个数,以及右侧 0 的个数。 |
| 70 | + |
48 | 71 | <!-- tabs:start -->
|
49 | 72 |
|
50 | 73 | ### **Python3**
|
51 | 74 |
|
52 | 75 | <!-- 这里可写当前语言的特殊实现逻辑 -->
|
53 | 76 |
|
54 | 77 | ```python
|
55 |
| - |
| 78 | +class BinaryIndexedTree: |
| 79 | + def __init__(self, n): |
| 80 | + self.n = n |
| 81 | + self.c = [0] * (n + 1) |
| 82 | + |
| 83 | + @staticmethod |
| 84 | + def lowbit(x): |
| 85 | + return x & -x |
| 86 | + |
| 87 | + def update(self, x, delta): |
| 88 | + while x <= self.n: |
| 89 | + self.c[x] += delta |
| 90 | + x += BinaryIndexedTree.lowbit(x) |
| 91 | + |
| 92 | + def query(self, x): |
| 93 | + s = 0 |
| 94 | + while x > 0: |
| 95 | + s += self.c[x] |
| 96 | + x -= BinaryIndexedTree.lowbit(x) |
| 97 | + return s |
| 98 | + |
| 99 | + |
| 100 | +class Solution: |
| 101 | + def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int: |
| 102 | + pos = {v: i for i, v in enumerate(nums2, 1)} |
| 103 | + ans = 0 |
| 104 | + n = len(nums1) |
| 105 | + tree = BinaryIndexedTree(n) |
| 106 | + for num in nums1: |
| 107 | + p = pos[num] |
| 108 | + left = tree.query(p) |
| 109 | + right = n - p - (tree.query(n) - tree.query(p)) |
| 110 | + ans += left * right |
| 111 | + tree.update(p, 1) |
| 112 | + return ans |
56 | 113 | ```
|
57 | 114 |
|
58 | 115 | ### **Java**
|
59 | 116 |
|
60 | 117 | <!-- 这里可写当前语言的特殊实现逻辑 -->
|
61 | 118 |
|
62 | 119 | ```java
|
| 120 | +class Solution { |
| 121 | + public long goodTriplets(int[] nums1, int[] nums2) { |
| 122 | + int n = nums1.length; |
| 123 | + int[] pos = new int[n]; |
| 124 | + BinaryIndexedTree tree = new BinaryIndexedTree(n); |
| 125 | + for (int i = 0; i < n; ++i) { |
| 126 | + pos[nums2[i]] = i + 1; |
| 127 | + } |
| 128 | + long ans = 0; |
| 129 | + for (int num : nums1) { |
| 130 | + int p = pos[num]; |
| 131 | + long left = tree.query(p); |
| 132 | + long right = n - p - (tree.query(n) - tree.query(p)); |
| 133 | + ans += left * right; |
| 134 | + tree.update(p, 1); |
| 135 | + } |
| 136 | + return ans; |
| 137 | + } |
| 138 | +} |
| 139 | + |
| 140 | +class BinaryIndexedTree { |
| 141 | + private int n; |
| 142 | + private int[] c; |
| 143 | + |
| 144 | + public BinaryIndexedTree(int n) { |
| 145 | + this.n = n; |
| 146 | + c = new int[n + 1]; |
| 147 | + } |
| 148 | + |
| 149 | + public void update(int x, int delta) { |
| 150 | + while (x <= n) { |
| 151 | + c[x] += delta; |
| 152 | + x += lowbit(x); |
| 153 | + } |
| 154 | + } |
| 155 | + |
| 156 | + public int query(int x) { |
| 157 | + int s = 0; |
| 158 | + while (x > 0) { |
| 159 | + s += c[x]; |
| 160 | + x -= lowbit(x); |
| 161 | + } |
| 162 | + return s; |
| 163 | + } |
| 164 | + |
| 165 | + public static int lowbit(int x) { |
| 166 | + return x & -x; |
| 167 | + } |
| 168 | +} |
| 169 | +``` |
63 | 170 |
|
| 171 | +### **C++** |
| 172 | + |
| 173 | +```cpp |
| 174 | +class BinaryIndexedTree { |
| 175 | +public: |
| 176 | + int n; |
| 177 | + vector<int> c; |
| 178 | + |
| 179 | + BinaryIndexedTree(int _n): n(_n), c(_n + 1){} |
| 180 | + |
| 181 | + void update(int x, int delta) { |
| 182 | + while (x <= n) |
| 183 | + { |
| 184 | + c[x] += delta; |
| 185 | + x += lowbit(x); |
| 186 | + } |
| 187 | + } |
| 188 | + |
| 189 | + int query(int x) { |
| 190 | + int s = 0; |
| 191 | + while (x > 0) |
| 192 | + { |
| 193 | + s += c[x]; |
| 194 | + x -= lowbit(x); |
| 195 | + } |
| 196 | + return s; |
| 197 | + } |
| 198 | + |
| 199 | + int lowbit(int x) { |
| 200 | + return x & -x; |
| 201 | + } |
| 202 | +}; |
| 203 | + |
| 204 | +class Solution { |
| 205 | +public: |
| 206 | + long long goodTriplets(vector<int>& nums1, vector<int>& nums2) { |
| 207 | + int n = nums1.size(); |
| 208 | + vector<int> pos(n); |
| 209 | + for (int i = 0; i < n; ++i) pos[nums2[i]] = i + 1; |
| 210 | + BinaryIndexedTree* tree = new BinaryIndexedTree(n); |
| 211 | + long long ans = 0; |
| 212 | + for (int& num : nums1) |
| 213 | + { |
| 214 | + int p = pos[num]; |
| 215 | + int left = tree->query(p); |
| 216 | + int right = n - p - (tree->query(n) - tree->query(p)); |
| 217 | + ans += 1ll * left * right; |
| 218 | + tree->update(p, 1); |
| 219 | + } |
| 220 | + return ans; |
| 221 | + } |
| 222 | +}; |
| 223 | +``` |
| 224 | +
|
| 225 | +### **Go** |
| 226 | +
|
| 227 | +```go |
| 228 | +type BinaryIndexedTree struct { |
| 229 | + n int |
| 230 | + c []int |
| 231 | +} |
| 232 | +
|
| 233 | +func newBinaryIndexedTree(n int) *BinaryIndexedTree { |
| 234 | + c := make([]int, n+1) |
| 235 | + return &BinaryIndexedTree{n, c} |
| 236 | +} |
| 237 | +
|
| 238 | +func (this *BinaryIndexedTree) lowbit(x int) int { |
| 239 | + return x & -x |
| 240 | +} |
| 241 | +
|
| 242 | +func (this *BinaryIndexedTree) update(x, delta int) { |
| 243 | + for x <= this.n { |
| 244 | + this.c[x] += delta |
| 245 | + x += this.lowbit(x) |
| 246 | + } |
| 247 | +} |
| 248 | +
|
| 249 | +func (this *BinaryIndexedTree) query(x int) int { |
| 250 | + s := 0 |
| 251 | + for x > 0 { |
| 252 | + s += this.c[x] |
| 253 | + x -= this.lowbit(x) |
| 254 | + } |
| 255 | + return s |
| 256 | +} |
| 257 | +
|
| 258 | +func goodTriplets(nums1 []int, nums2 []int) int64 { |
| 259 | + n := len(nums1) |
| 260 | + pos := make([]int, n) |
| 261 | + for i, v := range nums2 { |
| 262 | + pos[v] = i + 1 |
| 263 | + } |
| 264 | + tree := newBinaryIndexedTree(n) |
| 265 | + var ans int64 |
| 266 | + for _, num := range nums1 { |
| 267 | + p := pos[num] |
| 268 | + left := tree.query(p) |
| 269 | + right := n - p - (tree.query(n) - tree.query(p)) |
| 270 | + ans += int64(left) * int64(right) |
| 271 | + tree.update(p, 1) |
| 272 | + } |
| 273 | + return ans |
| 274 | +} |
64 | 275 | ```
|
65 | 276 |
|
66 | 277 | ### **TypeScript**
|
67 | 278 |
|
68 | 279 | ```ts
|
69 |
| - |
70 | 280 | ```
|
71 | 281 |
|
72 | 282 | ### **...**
|
|
0 commit comments