刷题笔记-小K的区间与值和 | 豆包MarsCode AI刷题

152 阅读5分钟

www.marscode.cn/practice/vk…

给定一个长度为 nn 的数组 aa,定义数组的权值为数组中任意两个数的按位与的值之和。具体来说,对于数组中的每个连续子数组,我们需要计算该子数组中所有可能的两个元素的按位与值之和,并将这些值相加。要求计算该数组中所有可能的连续子数组的权值和,最终结果对 109+710^9 + 7 取模。

示例:

  • 输入:n = 4, a = [2, 3, 1, 2]
  • 输出:16

解题思路

这道题的关键在于高效地计算所有连续子数组内所有元素对的按位与值之和。直接的暴力解法需要 O(n3)O(n^3) 的时间复杂度,无法通过大型数据的测试。因此,我们需要找到一种优化的算法来降低时间复杂度。

观察与分析

  1. 按位与的性质:

    • 按位与运算的结果仅在对应位两个操作数同时为 1 时,该位结果才为 1。
    • 对于一组数,如果我们能知道在某一位上为 1 的数的分布情况,就能计算出这一位对最终结果的贡献。
  2. 连续子数组的特性:

    • 数组的所有连续子数组的总数为 O(n2)O(n^2),但我们需要考虑的是如何在 O(n)O(n)O(nlogn)O(n \log n) 的时间内计算出所有子数组中元素对的按位与之和。
  3. 按位分治思想:

    • 我们可以逐位考虑每一个二进制位,计算每一位对最终结果的贡献。
    • 对于数组中的每一个二进制位,我们只需要统计在所有连续子数组中,该位为 1 的元素对数。

算法设计

  1. 逐位计算:

    • 我们遍历整数的每一个二进制位(通常是 32 位)。
    • 对于第 bb 位,找到数组中第 bb 位为 1 的所有元素的位置。
  2. 统计每一位的贡献:

    • 对于第 bb 位为 1 的元素,我们需要计算这些元素在所有连续子数组中两两组合的次数。
    • 我们可以先找到第 bb 位为 1 的所有元素的下标,设为数组 idx_b
  3. 计算连续子数组中元素对的计数:

    • **组合计数:**对于 idx_b 中的下标对 (i,j)(i, j),我们需要计算在多少个连续子数组中,这两个下标的元素会被同时包含。
    • **子数组包含条件:**下标 iijj 被同时包含在某个连续子数组中的数量等于: [ (\text{左侧选择数}) \times (\text{右侧选择数}) = (i - L + 1) \times (R - j + 1) ] 其中 LLRR 分别是数组的起始和结束位置,对于整个数组,L=0L = 0R=n1R = n - 1
  4. 优化计算:

    • 为了避免三重循环,我们可以利用前缀和或计数的方法来优化计算过程。
    • 对于位为 1 的下标列表 idx_b,我们可以累计计算贡献,而不需要遍历每一个可能的子数组。

算法步骤

  1. 初始化结果 total_weight = 0,定义取模常数 mod = 10^9 + 7

  2. 遍历每一位 $b$(从 0 到 31):

    a. 找出数组中第 $b$ 位为 1 的所有下标,存入数组 idx_b

    b. 如果 idx_b 的长度小于 2,说明该位的贡献为 0,跳过。

    c. 初始化变量,用于计算该位的总贡献。

  3. 对于 idx_b,计算其贡献:

    a. 对于每一个下标 i(从 0 到 k - 2),

    b. 计算如下内容:

    • 左侧选择数(idx_b[i] - prev_idx + 1),其中 prev_idx 是前一个下标。
    • 右侧选择数(next_idx - idx_b[i]),其中 next_idx 是后一个下标。
    • 元素对数:左侧选择数乘以右侧选择数。
    • 将元素对数乘以该位的值 (1 << b),累加到当前位的贡献中。
  4. 将当前位的贡献累加到 total_weight 中,注意取模。

  5. 最终返回 total_weight

时间复杂度分析

  • 外层循环遍历二进制的每一位,共 32 次。
  • 内层对每一位为 1 的下标列表进行遍历,总的时间复杂度为 O(32n)O(32n)
  • 总体时间复杂度为 O(n)O(n)

代码实现

def solution(n: int, a: list) -> int:
    mod = 10**9 + 7
    total_weight = 0
    max_bit = max(a).bit_length()
    for b in range(max_bit):
        idx_b = []
        for idx in range(n):
            if (a[idx] >> b) & 1:
                idx_b.append(idx)
        k = len(idx_b)
        if k < 2:
            continue
        # 计算该位的贡献
        contrib = 0
        prev_idx = -1
        for i in range(k):
            idx = idx_b[i]
            left = idx - prev_idx
            if i + 1 < k:
                next_idx = idx_b[i + 1]
            else:
                next_idx = n
            right = next_idx - idx
            count = left * right - 1
            contrib = (contrib + count * (1 << b)) % mod
            prev_idx = idx
        total_weight = (total_weight + contrib) % mod
    return total_weight

if __name__ == '__main__':
    print(solution(4, [2, 3, 1, 2]) == 16)
    print(solution(3, [5, 6, 7]) == 25)
    print(solution(2, [1, 10]) == 0)
    print(solution(5, [1, 2, 4, 8, 16]) == 0)

总结

通过利用按位与的性质和连续子数组的组合特点,我们成功地将时间复杂度从 O(n3)O(n^3) 优化到了 O(n)O(n)。核心思想是逐位计算每一位对最终结果的贡献,避免了重复计算。

在算法设计中,充分利用位运算的性质和数组下标的关系,可以大大提高计算效率。这种位分治的思想在处理与位运算相关的问题时非常有用。

算法分析

  • 时间复杂度:O(n×logM)O(n \times \log M),其中 MM 是数组中最大元素,logM\log M 表示其二进制位数,通常为 32。

  • 空间复杂度:O(n)O(n),用于存储每一位为 1 的下标列表。

可能的优化

  • 如果数组中的数很大,但实际有效位很少,可以通过只遍历数组中出现的有效位,进一步减少不必要的计算。

  • 在计算左右选择数时,可以预先计算并存储,以避免重复计算。

体会

这道题考查了对位运算性质的理解以及如何将暴力算法优化为线性算法的能力。在面对可能出现高时间复杂度的问题时,尝试从数学性质和算法优化的角度入手,往往可以找到突破口。