Python数组计算优化：解决大规模查询超时问题

任务描述 #

给定一个包含0 < N ≤ 100000个元素的数组a，以及0 < M ≤ 100000个查询。查询分为两种类型：

• ? l r k —— 统计数组区间[l, r)中值等于k的元素数量
• + i delta —— 将数组第i个元素的值增加delta

保证数组元素始终满足0 ≤ a[i] ≤ K。

输入格式 #

1. 第一行包含3个整数N、M、K，其中0 < N, M ≤ 100000，0 < K ≤ 100。N为数组元素个数，M为查询次数，K为数组元素的上限值。
2. 第二行包含N个整数，为数组a的初始值，满足0 ≤ a[i] ≤ K。
3. 接下来M行，每行包含一个字符及3或4个整数，代表两种查询类型之一：
   • ? l r k：0 ≤ l < r ≤ N，0 ≤ k ≤ K
   • + i delta：0 ≤ i < N，-K ≤ delta ≤ K

输出格式 #

对于每个?类型的查询，输出对应的统计结果。

输入样例 #

2 8 1
0 0
? 0 2 0
? 0 2 1
+ 0 1
? 0 2 0
? 0 2 1
+ 1 1
? 0 2 0
? 0 2 1

输出样例 #

2
0
1
1
0
2

现有代码问题分析 #

提供的代码中，FenwickTree类并未实现真正的树状数组结构，sum方法直接对数组切片调用count，时间复杂度为O(r-l)。当查询次数达到1e5且每次查询覆盖大部分数组时，总时间复杂度会达到O(M*N)，远超时间限制，导致超时。

优化方案与代码实现 #

由于K的取值上限仅为100，我们可以为每个数值0~K分别维护一个树状数组（Fenwick Tree），每个树状数组记录对应数值在数组各位置的出现次数：

• 更新操作：当数组第i个元素从old_val变为new_val时，在old_val对应的树状数组中对位置i+1（树状数组通常从1开始索引）减1，在new_val对应的树状数组中对位置i+1加1，同时更新原数组的值。
• 查询操作：对于? l r k，查询k对应的树状数组中[1, r]的前缀和减去[1, l]的前缀和，结果即为区间[l, r)中值为k的元素数量。

这种方案下，单次更新和查询的时间复杂度均为O(logN)，总时间复杂度为O((N+M)KlogN)，完全满足题目要求。

优化后的代码如下：

import sys

class FenwickTree:
    def __init__(self, size):
        self.n = size
        self.tree = [0]*(self.n + 1)  # 树状数组从1开始索引
    
    def update(self, idx, delta):
        # idx是1-based的位置
        while idx <= self.n:
            self.tree[idx] += delta
            idx += idx & -idx
    
    def query(self, idx):
        # 查询[1, idx]的前缀和（1-based）
        res = 0
        while idx > 0:
            res += self.tree[idx]
            idx -= idx & -idx
        return res

if __name__ == '__main__':
    input = sys.stdin.read().split()
    ptr = 0
    n = int(input[ptr])
    ptr +=1
    m = int(input[ptr])
    ptr +=1
    k = int(input[ptr])
    ptr +=1
    
    a = list(map(int, input[ptr:ptr+n]))
    ptr +=n
    
    # 为每个0~k的数值创建树状数组
    trees = [FenwickTree(n) for _ in range(k+1)]
    for i in range(n):
        val = a[i]
        trees[val].update(i+1, 1)  # 转换为1-based索引
    
    for _ in range(m):
        op = input[ptr]
        ptr +=1
        if op == '?':
            l = int(input[ptr])
            ptr +=1
            r = int(input[ptr])
            ptr +=1
            k_val = int(input[ptr])
            ptr +=1
            # 查询[1, r] - [1, l]，对应原数组[l, r)
            cnt = trees[k_val].query(r) - trees[k_val].query(l)
            print(cnt)
        else:
            i = int(input[ptr])
            ptr +=1
            delta = int(input[ptr])
            ptr +=1
            old_val = a[i]
            new_val = old_val + delta
            # 更新两个树状数组
            trees[old_val].update(i+1, -1)
            trees[new_val].update(i+1, 1)
            a[i] = new_val

代码说明 #

1. 树状数组实现：FenwickTree类提供update（单点更新）和query（前缀和查询）方法，均为O(logN)时间复杂度。
2. 批量读取输入：使用sys.stdin.read()一次性读取所有输入，避免多次input()调用的IO开销，提升速度。
3. 多树状数组维护：针对每个可能的数值（0到K）维护独立的树状数组，确保更新和查询的高效性。

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