查找从当前索引到所有其他索引的距离总和

Question

例如给定一个数组arr(0, 0, 1, 1, 1)，如果值为1，我需要找到所有其他索引到当前索引的距离之和。例如，我在O(n^2)中有以下代码表示法，但我知道我可能可以在 O(n) 中完成。

arr = [0, 1, 1, 1, 1]
for(i = 0; i < arr.length; i++) {
   sum = 0;
   for(j = 0; j < arr.length; j++) {
      if(arr[j] == 1 && j != i) {
         sum += Math.abs(j - i);
      }
   }
   return_arr[i] = sum;
}

（没有任何特定语言）

距离用j-i的绝对值表示。

Answer 1

无论该值是否为 1，任何索引 i 左侧的值数为 i，其右侧的值数为 n-1-i，其中 n 是数组的长度。

你知道前 n 个自然数的总和，即 1, 2, 3, ... n 是 n(n+1)/2

你对左边的索引数求和 = i(i+1)/2

右侧索引数的总和 = (n-i-1)(n-i)/2

将两个总和相加得到一个索引的总和 = i(i+1)/2 + (n-i-1)(n-i)/2

对每个值为 1 的索引执行此操作。

伪代码：

GetSumAtEachIndex(arr) {
    n = arr.length
    sum = array(n)
    for ( i = 0; i < arr.length; i++ ) {
        if ( arr[i] == 1 ) {
           sum[i] = i*(i+1)/2 + (n-i-1)*(n-i)/2;
        }
    }
    return sum;
}