数组子数组按位或总和计算超时，如何优化时间复杂度？

问题说明 #

给定长度为N的整数数组A，子数组的值定义为该子数组内所有元素的按位或结果，要求返回所有子数组的值的总和对10^9 + 7取模的结果。

示例 #

输入A = [1, 2, 3, 4, 5]，所有子数组的按位或值计算如下：

Value([1]) = 1
Value([1, 2]) = 3
Value([1, 2, 3]) = 3
Value([1, 2, 3, 4]) = 7
Value([1, 2, 3, 4, 5]) = 7
Value([2]) = 2
Value([2, 3]) = 3
Value([2, 3, 4]) = 7
Value([2, 3, 4, 5]) = 7
Value([3]) = 3
Value([3, 4]) = 7
Value([3, 4, 5]) = 7
Value([4]) = 4
Value([4, 5]) = 5
Value([5]) = 5
Sum of all these values = 71

已尝试的解法 #

1. 暴力解法 #

最直接的暴力枚举所有子数组计算按位或，时间复杂度O(n²)，小数据量正确，大数据量超时：

def solve(a):
        n = len(a)
        mod = 10 ** 9 + 7
        sum2 = 0
        for i in range(n):
            sum1 = 0
            for j in range(i, n):
                sum1 = (sum1 | a[j])
                sum2 = sum2 + sum1
        return sum2 % mod

2. 按位统计贡献解法 #

核心思路：每一个为1的比特位的总贡献 = 该位在所有子数组按位或结果中出现的次数 × 2的对应次方。
比如示例[5, 9, 14]的计算过程：
所有子数组的按位或结果：

• OR of [5] = 5 (0101)
• OR of [9] = 9 (1001)
• OR of [14] = 14 (1110)
• OR of [5,9] = 15 (1111)
• OR of [9,14] = 15 (1111)
• OR of [5,9,14] = 15 (1111)
  各比特位贡献计算：
• 第3位为1的次数是5，贡献为5×2^3
• 第2位为1的次数是5，贡献为5×2^2
• 第1位为1的次数是4，贡献为4×2^1
• 第0位为1的次数是5，贡献为5×2^0
  总贡献：5×8 + 5×4 + 4×2 + 5×1 = 71
  基于该思路编写的代码如下，但是提交仍然超时：

def solve(a):
        mod = 10 ** 9 + 7
        n = len(a)
        ans = 0
        for i in range(32):
            ind = n
            for j in range(n-1, -1,-1):
                bit = bool((a[j] >> i) &1)
                if bit:
                    ind = j
                ans += (n-ind)*pow(2,i)  
    
        return ans % mod   
solve([1,2,3,4,5])

优化方案 #

按位统计的思路时间复杂度已经是O(32n)，属于该问题的最优时间复杂度级别，出现超时主要是代码实现上的冗余开销导致，优化点如下：

1. 提前预计算32位对应的2的幂值，避免每次循环重复调用pow(2,i)计算幂次
2. 每一位统计完所有出现次数后再统一计算贡献，减少累加过程中的重复乘法操作
3. 去掉不必要的bool()类型转换，直接用位运算结果判断，提升计算速度
   优化后的代码如下：

def solve(a):
    mod = 10**9 +7
    n = len(a)
    ans = 0
    # 预计算2的0~31次方
    pow2 = [1 << i for i in range(32)]
    for bit in range(32):
        last = n
        cnt = 0
        for j in range(n-1, -1, -1):
            if (a[j] >> bit) & 1:
                last = j
            cnt += n - last
        # 按位累加贡献
        ans = (ans + cnt * pow2[bit]) % mod
    return ans

如果还是需要进一步压缩运行时间，也可以采用维护以当前元素结尾的所有不同或值的思路，利用按位或的单调性（每次新增元素或值只会增加不会减少，每个位置最多有32个不同的或值），同样是O(32n)的时间复杂度，常数级开销更低，实现如下：

def solve(a):
    mod = 10**9 +7
    ans = 0
    # 存储(or值, 出现次数)
    prev = []
    for num in a:
        curr = [(num, 1)]
        for val, cnt in prev:
            new_val = val | num
            if curr[-1][0] == new_val:
                curr[-1] = (new_val, curr[-1][1] + cnt)
            else:
                curr.append((new_val, cnt))
        # 累加当前位置的所有贡献
        for val, cnt in curr:
            ans = (ans + val * cnt) % mod
        prev = curr
    return ans

内容的提问来源于stack exchange，提问作者abd