Medium

题目描述

给你一个下标从 0 开始的整数数组 nums。现有一个长度等于 nums.length 的数组 arr，其中 arr[i] 是 |i - j| 对所有满足 nums[j] == nums[i] 且 j != i 的 j 求和。如果不存在这样的 j，则令 arr[i] = 0。

返回数组 arr。

示例 1：

输入：nums = [1,3,1,1,2]
输出：[5,0,3,4,0]
解释：
当 i = 0 时，nums[0] == nums[2] 且 nums[0] == nums[3]。因此，arr[0] = |0 - 2| + |0 - 3| = 5。
当 i = 1 时，arr[1] = 0，因为不存在值为 3 的其他下标。
当 i = 2 时，nums[2] == nums[0] 且 nums[2] == nums[3]。因此，arr[2] = |2 - 0| + |2 - 3| = 3。
当 i = 3 时，nums[3] == nums[0] 且 nums[3] == nums[2]。因此，arr[3] = |3 - 0| + |3 - 2| = 4。
当 i = 4 时，arr[4] = 0，因为不存在值为 2 的其他下标。

示例 2：

输入：nums = [0,5,3]
输出：[0,0,0]
解释：由于 nums 中每个元素都不相同，对于所有 i，都有 arr[i] = 0。

提示：

  • 1 <= nums.length <= 10^5
  • 0 <= nums[i] <= 10^9

解题思路

这道题需要计算每个位置到所有相同值位置的距离和。直接暴力计算会超时，需要使用前缀和优化。

核心思路：

  1. 首先用哈希表收集每个值出现的所有位置索引
  2. 对于每个值，其所有位置索引构成一个有序数组
  3. 对于位置 i，需要计算它到左边所有相同值位置和右边所有相同值位置的距离和

前缀和优化：

  • 对于位置 i，设左边有 left_count 个相同值位置，右边有 right_count 个相同值位置
  • 左边距离和 = i × left_count - 左边位置和
  • 右边距离和 = 右边位置和 - i × right_count
  • 使用前缀和可以快速计算位置和

具体步骤：

  1. 遍历数组，用哈希表记录每个值的所有位置
  2. 对每个值的位置数组计算前缀和
  3. 对于每个位置，根据它在该值位置数组中的索引，利用前缀和公式计算距离和

时间复杂度 O(n)，空间复杂度 O(n)，是最优解法。

代码实现

class Solution {
public:
    vector<long long> distance(vector<int>& nums) {
        unordered_map<int, vector<int>> indices;
        
        // 收集每个值的所有位置
        for (int i = 0; i < nums.size(); i++) {
            indices[nums[i]].push_back(i);
        }
        
        vector<long long> result(nums.size(), 0);
        
        // 对每个值计算距离和
        for (auto& [val, pos] : indices) {
            if (pos.size() <= 1) continue;
            
            // 计算前缀和
            vector<long long> prefix(pos.size() + 1, 0);
            for (int i = 0; i < pos.size(); i++) {
                prefix[i + 1] = prefix[i] + pos[i];
            }
            
            // 计算每个位置的距离和
            for (int i = 0; i < pos.size(); i++) {
                int curr = pos[i];
                long long leftSum = prefix[i];
                long long rightSum = prefix[pos.size()] - prefix[i + 1];
                
                result[curr] = (long long)curr * i - leftSum + 
                              rightSum - (long long)curr * (pos.size() - i - 1);
            }
        }
        
        return result;
    }
};

class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        from collections import defaultdict
        
        indices = defaultdict(list)
        
        # 收集每个值的所有位置
        for i, num in enumerate(nums):
            indices[num].append(i)
        
        result = [0] * len(nums)
        
        # 对每个值计算距离和
        for val, pos in indices.items():
            if len(pos) <= 1:
                continue
            
            # 计算前缀和
            prefix = [0] * (len(pos) + 1)
            for i in range(len(pos)):
                prefix[i + 1] = prefix[i] + pos[i]
            
            # 计算每个位置的距离和
            for i in range(len(pos)):
                curr = pos[i]
                left_sum = prefix[i]
                right_sum = prefix[len(pos)] - prefix[i + 1]
                
                result[curr] = curr * i - left_sum + right_sum - curr * (len(pos) - i - 1)
        
        return result

public class Solution {
    public long[] Distance(int[] nums) {
        Dictionary<int, List<int>> indices = new Dictionary<int, List<int>>();
        
        // 收集每个值的所有位置
        for (int i = 0; i < nums.Length; i++) {
            if (!indices.ContainsKey(nums[i])) {
                indices[nums[i]] = new List<int>();
            }
            indices[nums[i]].Add(i);
        }
        
        long[] result = new long[nums.Length];
        
        // 对每个值计算距离和
        foreach (var kvp in indices) {
            List<int> pos = kvp.Value;
            if (pos.Count <= 1) continue;
            
            // 计算前缀和
            long[] prefix = new long[pos.Count + 1];
            for (int i = 0; i < pos.Count; i++) {
                prefix[i + 1] = prefix[i] + pos[i];
            }
            
            // 计算每个位置的距离和
            for (int i = 0; i < pos.Count; i++) {
                int curr = pos[i];
                long leftSum = prefix[i];
                long rightSum = prefix[pos.Count] - prefix[i + 1];
                
                result[curr] = (long)curr * i - leftSum + 
                              rightSum - (long)curr * (pos.Count - i - 1);
            }
        }
        
        return result;
    }
}

var distance = function(nums) {
    const indices = new Map();
    
    // 收集每个值的所有位置
    for (let i = 0; i < nums.length; i++) {
        if (!indices.has(nums[i])) {
            indices.set(nums[i], []);
        }
        indices.get(nums[i]).push(i);
    }
    
    const result = new Array(nums.length).fill(0);
    
    // 对每个值计算距离和
    for (const [val, pos] of indices) {
        if (pos.length <= 1) continue;
        
        // 计算前缀和
        const prefix = new Array(pos.length + 1).fill(0);
        for (let i = 0; i < pos.length; i++) {
            prefix[i + 1] = prefix[i] + pos[i];
        }
        
        // 计算每个位置的距离和
        for (let i = 0; i < pos.length; i++) {
            const curr = pos[i];
            const leftSum = prefix[i];
            const rightSum = prefix[pos.length] - prefix[i + 1];
            
            result[curr] = curr * i - leftSum + 
                          rightSum - curr * (pos.length - i - 1);
        }
    }
    
    return result;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n)遍历数组一次收集位置，然后对每个位置计算距离和，总体为线性时间
空间复杂度O(n)使用哈希表存储位置信息和前缀和数组

相关题目