class Solution {
public:
vector<long long> distance(vector<int>& nums) {
}
};
6360. 等值距离和
给你一个下标从 0 开始的整数数组 nums 。现有一个长度等于 nums.length 的数组 arr 。对于满足 nums[j] == nums[i] 且 j != i 的所有 j ，arr[i] 等于所有 |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 <= 1050 <= nums[i] <= 109原站题解
python3 解法, 执行用时: 436 ms, 内存消耗: 50.9 MB, 提交时间: 2023-04-11 09:44:18
class Solution:
def distance(self, nums: List[int]) -> List[int]:
groups = defaultdict(list)
for i, x in enumerate(nums):
groups[x].append(i) # 相同元素分到同一组，记录下标
ans = [0] * len(nums)
for a in groups.values():
n = len(a)
s = sum(x - a[0] for x in a) # a[0] 到其它下标的距离之和
ans[a[0]] = s
for i in range(1, n):
# 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
s += (i * 2 - n) * (a[i] - a[i - 1])
ans[a[i]] = s
return ans
java 解法, 执行用时: 28 ms, 内存消耗: 73.6 MB, 提交时间: 2023-04-11 09:44:05
class Solution {
public long[] distance(int[] nums) {
int n = nums.length;
var groups = new HashMap<Integer, List<Integer>>();
for (int i = 0; i < n; ++i) // 相同元素分到同一组，记录下标
groups.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
var ans = new long[n];
for (var a : groups.values()) {
int m = a.size();
long s = 0;
for (int x : a)
s += x - a.get(0); // a[0] 到其它下标的距离之和
ans[a.get(0)] = s;
for (int i = 1; i < m; ++i)
// 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
ans[a.get(i)] = s += (long) (i * 2 - m) * (a.get(i) - a.get(i - 1));
}
return ans;
}
}
golang 解法, 执行用时: 104 ms, 内存消耗: 22.2 MB, 提交时间: 2023-04-11 09:43:46
// 相同元素分组+考虑增量
func distance(nums []int) []int64 {
groups := map[int][]int{}
for i, x := range nums {
groups[x] = append(groups[x], i) // 相同元素分到同一组，记录下标
}
ans := make([]int64, len(nums))
for _, a := range groups {
n := len(a)
s := int64(0)
for _, x := range a {
s += int64(x - a[0]) // a[0] 到其它下标的距离之和
}
ans[a[0]] = s
for i := 1; i < n; i++ {
// 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
s += int64(i*2-n) * int64(a[i]-a[i-1])
ans[a[i]] = s
}
}
return ans
}
golang 解法, 执行用时: 144 ms, 内存消耗: 24.3 MB, 提交时间: 2023-04-11 09:43:02
func distance(nums []int) []int64 {
groups := map[int][]int{}
for i, x := range nums {
groups[x] = append(groups[x], i) // 相同元素分到同一组，记录下标
}
ans := make([]int64, len(nums))
for _, a := range groups {
n := len(a)
s := make([]int, n+1)
for i, x := range a {
s[i+1] = s[i] + x // 前缀和
}
for i, target := range a {
left := target*i - s[i] // 蓝色面积
right := s[n] - s[i] - target*(n-i) // 绿色面积
ans[target] = int64(left + right)
}
}
return ans
}
java 解法, 执行用时: 20 ms, 内存消耗: 61.5 MB, 提交时间: 2023-04-11 09:42:41
class Solution {
public long[] distance(int[] nums) {
int n = nums.length;
var groups = new HashMap<Integer, List<Integer>>();
for (int i = 0; i < n; ++i) // 相同元素分到同一组，记录下标
groups.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
var ans = new long[n];
var s = new long[n + 1];
for (var a : groups.values()) {
int m = a.size();
for (int i = 0; i < m; ++i)
s[i + 1] = s[i] + a.get(i); // 前缀和
for (int i = 0; i < m; ++i) {
int target = a.get(i);
long left = (long) target * i - s[i]; // 蓝色面积
long right = s[m] - s[i] - (long) target * (m - i); // 绿色面积
ans[target] = left + right;
}
}
return ans;
}
}
python3 解法, 执行用时: 448 ms, 内存消耗: 51.3 MB, 提交时间: 2023-04-11 09:41:54
'''
相同元素分组+前缀和+二分查找
'''
class Solution:
def distance(self, nums: List[int]) -> List[int]:
groups = defaultdict(list)
for i, x in enumerate(nums):
groups[x].append(i) # 相同元素分到同一组，记录下标
ans = [0] * len(nums)
for a in groups.values():
n = len(a)
s = list(accumulate(a, initial=0)) # 前缀和
for j, target in enumerate(a):
left = target * j - s[j] # 蓝色面积
right = s[n] - s[j] - target * (n - j) # 绿色面积
ans[target] = left + right
return ans