304计算特定条件下的四元组数量

题解

该题目要求在给定数组中找到满足条件 a[i] + a[j] = a[k] ^ a[l] 且 i < j < k < l 的四元组数量。由于答案可能非常大，需要对 10^9 + 7 取模。

解题思路

1. 枚举所有可能的 (i, j) 对：

   • 对于每个 数对a[i]和a[j] (i < j)，计算 a[i] + a[j] 的值，并将其存储在字典 d 中，并记录对应a[j]的下标
   • 对字典 d 中每个键对应的列表进行排序，以便后续使用二分查找。

2. 枚举所有可能的 (k, l) 对：

   • 对于每个 (k, l)，满足k < l，计算 a[k] ^ a[l] 的值。寻找是否存在 (i, j) 满足 a[i] + a[j] = a[k] ^ a[l] 且 j < k。
   • 使用 bisect_left 在 d[a[k] ^ a[l]] 中找到满足 j < k 的 j 的数量，并累加到结果 res 中。

代码

mod = int(1e9 + 7)
from collections import defaultdict
from bisect import bisect_left
def solution(n: int, a: list) -> int:
    d = defaultdict(list)
    for i in range(n):
        for j in range(i + 1, n):
            d[a[i] + a[j]].append(j)
    for k in d:
        d[k].sort()
    res = sum(bisect_left(d[a[i] ^ a[j]], i) for i in range(2, n - 1) for j in range(i + 1, n))
    return res % mod

if __name__ == '__main__':
    print(solution(5, [2, 3, 1, 5, 4]) == 1)
    print(solution(6, [1, 2, 3, 4, 5, 6]) == 1)
    print(solution(4, [4, 1, 3, 2]) == 0)
    print(solution(14, [9,2,2,5,6,9,5,7,13,9,15,16,5,10]) == 33)

复杂度分析

• 时间复杂度：O(n^2logn)
  • 枚举 (i, j) 对需要 O(n^2) 时间。
  • 对每个 (k, l) 对进行二分查找需要 O(log n) 时间，总体为 O(n^2logn)。
• 空间复杂度：O(n^2)
  • 用于存储所有 (i, j) 对的和。