Problem: 6360. 等值距离和

[TOC]

思路

因为肯定要找出 nums[i]=nums[j] 的所有 i,j，所以首先遍历一遍，记录各个数值出现的每个下标。拿到了各个数值的下标数组后，对数组中的各个元素，求它与其他元素之间的距离和。

关键在于，如何快速地求出数组中元素与其他元素的距离和？在我不知道“前缀和”的概念及其运用方法之前，我是直接遍历求和的，由于题目说 nums.length <= 10⁵, 不出意外地超时了。随后我从题目提示中得知了“前缀和”的概念，并从其他题解中了解了前缀和在本题中的运用，故写此题解给像我一样的小白讲解前缀和与本题的关系。

解题方法

前缀和

参考 https://www.cnblogs.com/ZghzzZyu/p/16407303.html

前缀和指的是，对于数组 nums, 有前缀和数组 arr, 其每一项 arr[i] 的值为 nums[0] + nums[1] + ... + nums[i] 的和。

$arr[i] = sum(nums[0...i])$

例如有数组 nums 如下：

$nums = [0, 1, 2, 3, 4]$

则对应的前缀和数组为：

$arr = [0, 1, 3, 6, 10]$

根据定义不难看出，前缀和数组的每一项，都为其前一项加上 nums 中当前的项。

$arr[i] = arr[i - 1] + nums[i]$

这使得我们在遍历 nums 的每一项时，可以立即计算出 arr 中对应的项，这也是利用前缀和解这道题时降低时间复杂度的关键。在思路中讲到，我们首先遍历 nums 并记录每个数字出现的下标数组，现在，我们改为在遍历 nums 时记录每个数字出现的下标数组的前缀和数组。

通过前缀和求距离和

参考 https://leetcode.cn/problems/sum-of-distances/solutions/2216328/an-shu-zi-xia-biao-qian-zhui-he-by-yusha-xzu4/

有了前缀和数组，就能很方便地求出距离和了。例如，如果我们拥有数组 nums 如下：

$nums = [0, 1, 2, 3, 4]$

当我们要求元素 $2$ 到其他元素的距离和时，计算过程如下：

对于小于 2 的元素 n，计算 $2 - n$ 并相加。不难发现，其结果为小于 2 的元素总和（即前缀和），减去 2 乘以小于 2 的元素数量

$(2 - 0) + (2 - 1) = 2*2 - (0 + 1) = 2 * count(0...1) - sum(0...1)$

对于大于 2 的元素 n，计算 $n - 2$ 并相加。不难发现，其结果为大于 2 的元素总和，减去 2 乘以大于 2 的元素数量

$(3 - 2) + (4 - 2) = (3 + 4) - 2 * 2 = sum(3...4) - 2 * count(3...4)$

根据前缀和的特性，$sum(3...4)$ 可以通过 $sum(0...4)$ 减去 $sum(0...2)$ 得出。

$sum(3...4) = sum(0...4) - sum(0...2)$

元素 2 到其他元素的距离和即为上面两部分相加。

Code

###TypeScript

function distance(nums: number[]): number[] {
  const answer = new Array<number>(nums.length).fill(0)
  // 用于记录各个数值出现的下标数组的前缀和数组
  const map = new Map<number, number[]>();

  for (let i = 0; i < nums.length; i++) {
    if (!map.has(nums[i])) {
      map.set(nums[i], [i])
    } else {
      const array = map.get(nums[i])
      // 计算前缀和
      array.push(array[array.length - 1] + i)
    }
  }

  // 用于记录当前下标是下标数组中的第几个元素
  const indexCounter = []
  for (let i = 0; i < nums.length; i++) {
    const array = map.get(nums[i])

    if (array.length > 1) {
      const indexOfIndex = typeof indexCounter[nums[i]] === 'undefined' ? (indexCounter[nums[i]] = 0) : (indexCounter[nums[i]])
      indexCounter[nums[i]]++

      // 套上面的计算公式
      answer[i] =
        (i * indexOfIndex - (array[indexOfIndex - 1] ?? 0)) +
        (array[array.length - 1] - array[indexOfIndex] - i * (array.length - 1 - indexOfIndex))
    }
  }

  return answer
};