题目地址:
https://www.lintcode.com/problem/distance-metrics/description
给定一个长 n n n的数组 A A A,其中可能有重复数字,要求返回一个数组 B B B,与 A A A同长,并且要求 B [ i ] = ∑ A [ k ] = A [ i ] ∣ k − i ∣ B[i]=\sum_{A[k]=A[i]}|k-i| B[i]=A[k]=A[i]∑∣k−i∣
思路是前缀和。首先开个哈希表,存一下每个数字出现在了哪些位置。不妨设 x x x出现的位置的数组是 c c c,也就是说 ∀ j , A [ c [ j ] ] = x \forall j,A[c[j]]=x ∀j,A[c[j]]=x。接着我们要算的 B [ c [ j ] ] B[c[j]] B[c[j]]满足: B [ c [ j ] ] = ∑ i < j ( c [ j ] − c [ i ] ) + ∑ j < k ( c [ k ] − c [ j ] ) = j c [ j ] − ∑ i = 0 j − 1 c [ i ] + ∑ k = j + 1 l c − 1 c [ k ] − ( l c − 1 − j ) c [ j ] B[c[j]]=\sum_{i<j}(c[j]-c[i])+\sum_{j<k}(c[k]-c[j])\\=jc[j]-\sum_{i=0}^{j-1}c[i]+\sum_{k=j+1}^{l_c-1}c[k]-(l_c-1-j)c[j] B[c[j]]=i<j∑(c[j]−c[i])+j<k∑(c[k]−c[j])=jc[j]−i=0∑j−1c[i]+k=j+1∑lc−1c[k]−(lc−1−j)c[j]由此看出,只需求出 c c c的前缀和就能加快计算。设 c c c的前缀和数组是 d d d,满足 d [ i ] = ∑ k < i c [ k ] d[i]=\sum_{k<i}c[k] d[i]=∑k<ic[k],那么有: B [ c [ j ] ] = j c [ j ] − d [ j ] + d [ l c ] − d [ j ] − ( l c − j ) c [ j ] B[c[j]]=jc[j]-d[j]+d[l_c]-d[j]-(l_c-j)c[j] B[c[j]]=jc[j]−d[j]+d[lc]−d[j]−(lc−j)c[j]写代码的时候注意要用long类型来写。代码如下:
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class Solution {
/**
* @param a: an integer array
* @return: an integer array
*/
public List<Long> getDistanceMetrics(int[] a) {
// write your code here.
List<Long> res = new ArrayList<>();
// 先填n个数
for (int i = 0; i < a.length; i++) {
res.add(0L);
}
// 开个哈希表存每个数出现在了哪些位置
Map<Integer, List<Integer>> map = new HashMap<>();
for (int i = 0; i < a.length; i++) {
map.putIfAbsent(a[i], new ArrayList<>());
map.get(a[i]).add(i);
}
// 开个哈希表存每个数出现的位置的数组的前缀和,方便做累加
Map<Integer, long[]> preSums = new HashMap<>();
for (Map.Entry<Integer, List<Integer>> entry : map.entrySet()) {
// 得到key出现的位置
List<Integer> list = entry.getValue();
// 如果只出现了一次,那么那个位置应该填0
if (list.size() == 1) {
res.set(list.get(0), 0L);
continue;
}
// 否则先求出位置数组的前缀和
long[] preSum = null;
if (preSums.containsKey(entry.getKey())) {
preSum = preSums.get(entry.getKey());
} else {
preSum = new long[list.size() + 1];
for (int i = 0; i < list.size(); i++) {
preSum[i + 1] = preSum[i] + list.get(i);
}
}
// 将前缀和存一下
preSums.putIfAbsent(entry.getKey(), preSum);
// 按照公式算即可
for (int i = 0; i < list.size(); i++) {
int pos = list.get(i);
res.set(pos, ((long) pos * i - preSum[i]) +
(preSum[list.size()] - preSum[i] - (long) pos * (list.size() - i)));
}
}
return res;
}
}
时空复杂度 O ( n ) O(n) O(n)。