【Lintcode】327. Distance Metrics

题目地址：

https://www.lintcode.com/problem/distance-metrics/description

给定一个长$n$的数组$A$，其中可能有重复数字，要求返回一个数组$B$，与$A$同长，并且要求$$B[i]=\sum _{A[k]=A[i]}|k-i|$$

思路是前缀和。首先开个哈希表，存一下每个数字出现在了哪些位置。不妨设$x$出现的位置的数组是$c$，也就是说$\forall j,A[c[j]]=x$。接着我们要算的$B[c[j]]$满足：$$\begin{gathered} 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] \end{gathered}$$由此看出，只需求出$c$的前缀和就能加快计算。设$c$的前缀和数组是$d$，满足$d[i]={\sum }_{k<i}c[k]$，那么有：$$B[c[j]]=jc[j]-d[j]+d[l_c]-d[j]-(l_c-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)$。