如何将O(N²)复杂度的0-1配对算法优化至更低复杂度以通过性能测试？

问题背景 #

需求是统计数组中有效0-1配对的数量：配对需满足「索引i < j，且a[i]=0、a[j]=1」，若结果超过1000000000则返回-1。原代码采用双重循环实现，正确性没问题，但时间复杂度为O(N²)，当数组规模超过10000时会超时，需要优化到线性复杂度。

原代码（Java）：

int result = 0;
int len = a.length;

for (int i = 0; i < len; i++) {
    int current = a[i];

    if (current != 0) {
        continue;
    }

    for (int inner = i; inner < len; inner++) {
        if (current != a[inner]) {
            result++;
        }
    }
}

return (result > 1000000000) ? -1 : result;

优化思路 #

原代码的核心问题是对每个0都要遍历后续所有元素统计1的数量，重复计算导致平方复杂度。换个角度优化：

1. 先统计数组中总共有多少个1
2. 遍历数组，遇到0时，将当前剩余的1的数量加到结果中（这个0可以和所有还没遍历到的1配对）
3. 遇到1时，将总1数减1（后续的0无法和这个1配对）
4. 每次累加后提前检查是否超过阈值，避免溢出

这种方法仅需遍历数组一次，时间复杂度O(N)，空间复杂度O(1)，完全能处理大规模输入。

优化后代码（Java） #

public int countValidPairs(int[] a) {
    final int MAX_LIMIT = 1000000000;
    int totalOnes = 0;
    // 先统计数组中1的总数
    for (int num : a) {
        if (num == 1) {
            totalOnes++;
        }
    }

    int result = 0;
    for (int num : a) {
        if (num == 0) {
            result += totalOnes;
            // 提前判断是否超出阈值，直接返回-1
            if (result > MAX_LIMIT) {
                return -1;
            }
        } else {
            totalOnes--;
        }
    }
    return result > MAX_LIMIT ? -1 : result;
}

示例验证 #

• 输入[0,1,0,1,1]：总1数初始为3，遍历过程中依次加3、加2，最终结果5，符合预期
• 输入[0,1,1,1,1]：总1数4，第一个0加4，后续1递减，最终结果4，符合预期
• 输入[0,1,0,0,1,0,1]：总1数3，依次加3、加2、加2、加1，最终结果8，符合预期

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