豆包MarsCode AI 刷题 中一道标记为困难的题目
题意如下所示:
小X有一个数组,她定义数组的权值为数组中任选两个数的按位或的值之和。具体来说,对于数组中的每个连续子数组,我们可以计算所有可能的两个元素的按位或值之和,并将这些值相加。小C想知道该数组中所有可能的连续子数组的权值和是多少,最后结果对10^9 + 7
取模。
题解
显然,可以通过暴力算法求解,时间复杂度为O(n^2),因为测试样例比较弱,所以也可以通过。
def solution(n: int, a: list) -> int:
ans = 0
mod = 10 ** 9 + 7
for i in range(1 , n + 1) :
x = a[i - 1]
for j in range(i + 1 , n + 1) :
y = a[j - 1]
ans += (x | y) * i * (n - j + 1) % mod
return ans
if __name__ == '__main__':
print(solution(4, [2, 3, 1, 2]) == 44)
print(solution(3, [5, 6, 7]) == 35)
print(solution(2, [1, 10]) == 11)
print(solution(5, [1, 2, 4, 8, 16]) == 402)
接下来,就考虑是否有时间复杂度可以降到nlogn的算法。我请教了AI,以下是AI的解答。
算法思路概述:
- 按位分离: 我们可以将问题分解为对于每一个二进制位,计算其对最终答案的贡献。
- 统计每个位的贡献: 对于每个二进制位,我们需要计算该位在所有可能的元素对中被设置为1的总次数,考虑到元素对在所有连续子数组中的重复出现,我们需要计算每个元素对在多少个连续子数组中出现。
- 优化计算: 利用栈或线段树等数据结构来高效地统计每个位的贡献,以满足 O(nlogn)的时间复杂度要求。
详细算法步骤:
-
初始化答案变量:
复制代码 MOD = 10**9 + 7 total_sum = 0
-
对于每个二进制位 bbb(从0到30,共31位):
复制代码 for b in range(31): # 初始化该位的贡献 bit_contribution = 0
-
预处理每个元素的第 bbb 位:
-
创建一个数组 S,其中 S[i]=1表示第 i个元素的第 b位为1,S[i]=0表示为0。
复制代码 S = [ (A[i] >> b) & 1 for i in range(n) ]
-
-
计算第 bbb 位的贡献:
-
统计连续的0序列: 对于第 b 位为0 的位置,找到所有连续的0的区间。
-
计算这些区间内元素对的数量: 对于每个连续的0的区间,长度为 l,则该区间内元素对数为 \frac{l \times (l - 1)}{2}2l×(l−1)。
-
计算总的元素对数: 数组总元素对数为 n×(n−1)2。
-
计算第 b 位为1的元素对数:
复制代码 total_pairs = n * (n - 1) // 2 zero_pairs = sum( l * (l - 1) // 2 for each zero sequence of length l ) one_pairs = total_pairs - zero_pairs
-
考虑元素对在所有连续子数组中的重复次数:
-
元素对 (i,j) 在多少个连续子数组中出现?
- 对于位置 i 和 j(i<j),它们被包含在的连续子数组数为 count=i×(n−j+1)。
-
-
计算第 b 位的总贡献:
复制代码 bit_contribution = (one_pairs * 2**b) % MOD
-
-
累加到总答案:
复制代码 total_sum = (total_sum + bit_contribution) % MOD
-
输出答案:
复制代码 print(total_sum)
以上的思路基本正确,核心在于按位分解。但是细节上还是有些错误的。