计算特定条件下的四元组数量 | 豆包MarsCode AI刷题

题目描述

小S正在处理一个数组，目标是找到满足特定条件的四元组数量。给定一个长度为n的数组a，他需要计算有多少四元组$(i,j,k,l)$满足： a_i + a_j = a_k ⊕ a_l，其中 $\oplus$ 表示按位异或，并且满足索引的条件 i < j < k < l。

由于答案可能非常大，所以你需要对 $10^9 + 7$ 取模后输出结果。

测试样例

样例1：

输入：n = 5, a = [2, 3, 1, 5, 4] 输出：1

样例2：

输入：n = 6, a = [1, 2, 3, 4, 5, 6] 输出：1

样例3：

输入：n = 4, a = [4, 1, 3, 2] 输出：0

解题思路

1. 理解问题：

   • 我们需要找到满足条件的四元组 $(i, j, k, l)$，其中 $i < j < k < l$。
   • 条件是 $a_i + a_j = a_k \oplus a_l$。

2. 数据结构选择：

   • 由于我们需要遍历所有可能的四元组，使用四重循环是直接的方法。

3. 算法步骤：

   • 使用四重循环遍历所有可能的四元组 $(i, j, k, l)$。
   • 检查每个四元组是否满足 $a_i + a_j = a_k \oplus a_l$。
   • 如果满足条件，计数器加1。
   • 最后对结果取模 $10^9 + 7$。

代码实现

def solution(n: int, a: list) -> int:
    MOD = 10**9 + 7
    count = 0
    
    # 使用四重循环来遍历所有可能的四元组 (i, j, k, l)
    for i in range(n):
        for j in range(i + 1, n):
            for k in range(j + 1, n):
                for l in range(k + 1, n):
                    if a[i] + a[j] == a[k] ^ a[l]:
                        count += 1
    
    return count % 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)

总结

• 该问题通过四重循环遍历所有可能的四元组来解决。
• 时间复杂度为 $O(n^4)$，适用于较小的数组。
• 结果需要对 $10^9 + 7$ 取模以防止溢出。

进一步优化思路

1. 优化方向：

   • 当前算法的时间复杂度为 $O(n^4)$，对于较大的数组可能会导致性能问题。
   • 可以考虑使用哈希表来存储中间结果，减少重复计算。

2. 优化步骤：

   • 使用哈希表存储 $a_i + a_j$ 的值及其对应的索引对 $(i, j)$。
   • 遍历 $a_k$ 和 $a_l$，查找哈希表中是否存在 $a_k \oplus a_l$ 的值。
   • 如果存在，检查索引是否满足 $i < j < k < l$。

优化代码框架

def solution(n: int, a: list) -> int:
    MOD = 10**9 + 7
    count = 0
    hash_map = {}
    
    # 使用哈希表存储 a_i + a_j 的值及其对应的索引对 (i, j)
    for i in range(n):
        for j in range(i + 1, n):
            sum_ij = a[i] + a[j]
            if sum_ij in hash_map:
                hash_map[sum_ij].append((i, j))
            else:
                hash_map[sum_ij] = [(i, j)]
    
    # 遍历 a_k 和 a_l，查找哈希表中是否存在 a_k ⊕ a_l 的值
    for k in range(n):
        for l in range(k + 1, n):
            xor_kl = a[k] ^ a[l]
            if xor_kl in hash_map:
                for (i, j) in hash_map[xor_kl]:
                    if i < j < k < l:
                        count += 1
    
    return count % 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)

总结

• 通过使用哈希表存储中间结果，可以将时间复杂度降低到 $O(n^2)$。
• 这种优化方法适用于较大的数组，能够显著提高性能。