classSolution: defgoodTriplets(self, nums1: List[int], nums2: List[int]) -> int: n = len(nums1) p = [0] * n for i, x inenumerate(nums1): p[x] = i ans = 0 tree = [0] * (n + 1) for i inrange(1, n - 1): # 将 p[nums2[i - 1]] + 1 加入树状数组 j = p[nums2[i - 1]] + 1 while j <= n: tree[j] += 1 j += j & -j # 计算 less y, less = p[nums2[i]], 0 j = y while j: less += tree[j] j &= j - 1 ans += less * (n - 1 - y - (i - less)) return ans
classSolution { public: longlonggoodTriplets(vector<int> &nums1, vector<int> &nums2){ int n = nums1.size(); vector<int> p(n); for (int i = 0; i < n; ++i) p[nums1[i]] = i; longlong ans = 0; vector<int> tree(n + 1); for (int i = 1; i < n - 1; ++i) { for (int j = p[nums2[i - 1]] + 1; j <= n; j += j & -j) // 将 p[nums2[i-1]]+1 加入树状数组 ++tree[j]; int y = p[nums2[i]], less = 0; for (int j = y; j; j &= j - 1) // 计算 less less += tree[j]; ans += (long) less * (n - 1 - y - (i - less)); } return ans; } };
[sol1-Go]
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
funcgoodTriplets(nums1, nums2 []int) (ans int64) { n := len(nums1) p := make([]int, n) for i, v := range nums1 { p[v] = i } tree := make([]int, n+1) for i := 1; i < n-1; i++ { for j := p[nums2[i-1]] + 1; j <= n; j += j & -j { // 将 p[nums2[i-1]]+1 加入树状数组 tree[j]++ } y, less := p[nums2[i]], 0 for j := y; j > 0; j &= j - 1 { // 计算 less less += tree[j] } ans += int64(less) * int64(n-1-y-(i-less)) } return }
附 Python SortedList 做法:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
from sortedcontainers import SortedList
classSolution: defgoodTriplets(self, nums1: List[int], nums2: List[int]) -> int: n = len(nums1) p = [0] * n for i, x inenumerate(nums1): p[x] = i ans = 0 s = SortedList() for i inrange(1, n - 1): s.add(p[nums2[i - 1]]) y = p[nums2[i]] less = s.bisect_left(y) ans += less * (n - 1 - y - (i - less)) return ans