HackerRank带0替换的子数组求和查询解法超时求优化方案

原实现的核心问题 #

• 时间复杂度过高：采用双层循环结构，每收到一个查询都要完整遍历对应区间的所有元素，整体复杂度为O(Q*N)（Q为查询总数，N为主数组长度）。当数据规模较大时（如N、Q均达到1e5量级），运算量会突破时间限制。
• 存在边界逻辑bug：内层循环终止条件为j < r，会遗漏右边界位置的元素（比如查询区间为1_{3时，转换为0索引后区间为0}2，原代码仅遍历索引0、1，漏掉索引2的元素）。
• 存在冗余无效逻辑：定义了count变量但全程未参与实际计算，属于无用代码。

优化核心思路 #

利用双前缀和数组将每个查询的计算从遍历降级为常数时间计算，整体预处理仅需一次O(N)遍历：

• 第一份前缀和数组：存储主数组内所有非零元素的累计和，即原数组中0值不贡献累计值，其他元素按原值累加。这样任意区间[l, r]内的非零元素总和，可通过前缀和数组的差分快速算出。
• 第二份前缀和数组：存储主数组内累计出现的0值个数。这样任意区间[l, r]内的0值总数，可通过该数组的差分快速算出。
• 每个查询的最终求和公式：区间非零元素和 + 区间0值个数 * 查询指定的x
  也就是：(non_zero_prefix[r] - non_zero_prefix[l-1]) + (zero_count_prefix[r] - zero_count_prefix[l-1]) * x

优化后的C++实现 #

vector<long> findSum(vector<int> numbers, vector<vector<int>> queries) {
    int n = numbers.size();
    // 构建两个前缀和数组，下标从1开始，统一处理边界避免额外判断
    vector<long> non_zero_prefix(n + 1, 0);
    vector<long> zero_count_prefix(n + 1, 0);
    // 单次遍历完成前缀和预处理
    for (int i = 0; i < n; i++) {
        non_zero_prefix[i+1] = non_zero_prefix[i];
        zero_count_prefix[i+1] = zero_count_prefix[i];
        if (numbers[i] == 0) {
            zero_count_prefix[i+1] += 1;
        } else {
            non_zero_prefix[i+1] += numbers[i];
        }
    }
    vector<long> result;
    for (auto& q : queries) {
        // 原查询采用1索引，与前缀和下标天然对齐，无需额外偏移
        int l = q[0];
        int r = q[1];
        long x = q[2];
        // 直接通过前缀和差分计算当前查询结果
        long sum = non_zero_prefix[r] - non_zero_prefix[l-1] 
                 + (zero_count_prefix[r] - zero_count_prefix[l-1]) * x;
        result.push_back(sum);
    }
    return result;

• 优化后整体时间复杂度为O(N + Q)，无论N和Q的规模多大，都能满足时间限制要求。
• 同时修正了原实现中右边界元素遗漏的问题，移除了冗余的无效变量。

内容的提问来源于stack exchange，提问作者Hua Chi Quan