如何用单循环计算字符串中相同数字的位置总距离（O(n)复杂度）

需求：计算字符串中所有'1'的出现位置之间的总距离，要求时间复杂度O(n)，不能用嵌套循环。例如字符串"100101"中，'1'的位置为0、3、5，总距离是3（0到3）+5（0到5）+2（3到5）=10。

你的现有代码先收集所有'1'的位置，再用嵌套循环计算两两距离之和，时间复杂度是O(m²)（m为'1'的个数），当字符串很长且'1'数量较多时，效率会大幅下降。

优化思路 #

不用先收集所有位置再计算，而是遍历字符串时实时维护两个变量：

• count：已经遇到的'1'的个数
• sum_pos：已经遇到的'1'的位置总和

每遇到一个新的'1'（位置为pos），它和之前所有count个'1'的距离总和为 pos * count - sum_pos。把这个值加到总距离里，然后更新count和sum_pos即可。整个过程只需要一次遍历字符串，时间复杂度O(n)，空间复杂度O(1)（无需存储所有位置）。

优化后的代码 #

def pairs(s):
    total = 0
    count = 0
    sum_pos = 0
    for pos, char in enumerate(s):
        if char == "1":
            # 当前位置与之前所有1的距离总和
            total += pos * count - sum_pos
            count += 1
            sum_pos += pos
    return total

if __name__ == "__main__":
    print(pairs("100101"))  # 输出10，和原代码一致
    print(pairs("101"))     # 输出2，和原代码一致
    print(pairs("100100111001"))  # 输出71，和原代码一致

代码逻辑解释 #

以"100101"为例：

1. 第一个'1'在位置0：此时count=0，无距离可加；count变为1，sum_pos更新为0。
2. 第二个'1'在位置3：计算3*1 - 0=3，total变为3；count变为2，sum_pos更新为0+3=3。
3. 第三个'1'在位置5：计算5*2 -3=7，total变为3+7=10；count变为3，sum_pos更新为3+5=8。

最终得到正确的总距离10，完全符合O(n)的时间复杂度要求，空间效率也更高。

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