2615.等值距离和

链接：2615.等值距离和
难度：Medium
标签：数组、哈希表、前缀和
简介：给你一个下标从 0 开始的整数数组 nums 。现有一个长度等于 nums.length 的数组 arr 。对于满足 nums[j] == nums[i] 且 j != i 的所有 j ，arr[i] 等于所有 |i - j| 之和。如果不存在这样的 j ，则令 arr[i] 等于 0 。返回数组 arr 。

题解 1 - cpp

• 编辑时间：2023-04-09
• 执行用时：72ms
• 内存消耗：34.7MB
• 编程语言：cpp
• 解法介绍：收集所有相同数字的下标，做前缀和进行左右比较。

class Solution {
public:
    typedef long long ll;
    vector<ll> distance(vector<int>& nums) {
        int n = nums.size();
        vector<ll> arr(n, 0);
        unordered_map<int, vector<int>> m;
        for (int i = 0; i < n; i++) {
            m[nums[i]].push_back(i);
        }
        for (auto &item : m) {
            auto &list = item.second;
            ll sum = 0, left = 0;
            for (auto &v : list) sum += v - list[0];
            for (int i = 0; i < list.size(); i++) {
                arr[list[i]] = sum + left;
                if (i + 1 < list.size()) {
                    left += (list[i + 1] - list[i]) * (i + 1);
                    sum -= (list[i + 1] - list[i]) * (list.size() - 1 - i);
                }
            }
        }
        return arr;
    }
};

题解 2 - python

• 编辑时间：2023-04-09
• 执行用时：412ms
• 内存消耗：52.2MB
• 编程语言：python
• 解法介绍：同上。

class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        n = len(nums)
        arr = [0] * n
        m: dict[int, List[int]] = {}
        for i in range(n):
            if not nums[i] in m:
                m[nums[i]] = []
            m[nums[i]].append(i)
        for _, list in m.items():
            l, r = 0, 0
            for v in list:
                r += v - list[0]
            for i in range(len(list)):
                arr[list[i]] = r + l
                if i + 1 < len(list):
                    l += (list[i + 1] - list[i]) * (i + 1)
                    r -= (list[i + 1] - list[i]) * (len(list) - 1 - i)
        return arr

题解 3 - rust

• 编辑时间：2023-04-09
• 执行用时：68ms
• 内存消耗：12MB
• 编程语言：rust
• 解法介绍：同上。

impl Solution {
    pub fn distance(nums: Vec<i32>) -> Vec<i64> {
        use std::collections::HashMap;
        let n = nums.len();
        let mut arr: Vec<i64> = vec![0; n];
        let mut m = HashMap::<i32, Vec<usize>>::new();
        for i in 0..n {
            let item = m.entry(nums[i]).or_insert(Vec::new());
            item.push(i);
        }
        for (_, list) in m.into_iter() {
            let (mut r, mut l) = (0, 0);
            for v in &list {
                r += *v - list[0];
            }
            for i in 0..list.len() {
                arr[list[i]] = (r + l) as i64;
                if i + 1 < list.len() {
                    l += (list[i + 1] - list[i]) * (i + 1);
                    r -= (list[i + 1] - list[i]) * (list.len() - 1 - i);
                }
            }
        }
        arr
    }
}