2721 Sum Of Distances

Sum of Distances

You are given a 0-indexed integer array nums. There exists an array arr of length nums.length, where arr[i] is the sum of

i - j

over all j such that nums[j] == nums[i] and j != i. If there is no such j, set arr[i] to be 0.

Return the array *arr.*

 

Example 1:

**Input:** nums = [1,3,1,1,2]
**Output:** [5,0,3,4,0]
**Explanation:** 
When i = 0, nums[0] == nums[2] and nums[0] == nums[3]. Therefore, arr[0] = |0 - 2| + |0 - 3| = 5. 
When i = 1, arr[1] = 0 because there is no other index with value 3.
When i = 2, nums[2] == nums[0] and nums[2] == nums[3]. Therefore, arr[2] = |2 - 0| + |2 - 3| = 3. 
When i = 3, nums[3] == nums[0] and nums[3] == nums[2]. Therefore, arr[3] = |3 - 0| + |3 - 2| = 4. 
When i = 4, arr[4] = 0 because there is no other index with value 2. 

Example 2:

**Input:** nums = [0,5,3]
**Output:** [0,0,0]
**Explanation:** Since each element in nums is distinct, arr[i] = 0 for all i.

 

Constraints:

1 <= nums.length <= 105
0 <= nums[i] <= 109

 

Note: This question is the same as 2121: Intervals Between Identical Elements.

use std::collections::HashMap;

impl Solution {
    pub fn distance(nums: Vec<i32>) -> Vec<i64> {
        let n = nums.len();
        let mut groups: HashMap<i32, Vec<usize>> = HashMap::new();
        for (i, &v) in nums.iter().enumerate() {
            groups.entry(v).or_default().push(i);
        }
        let mut res = vec![0i64; n];
        for group in groups.values() {
            let total: i64 = group.iter().map(|&x| x as i64).sum();
            let mut prefix_total: i64 = 0;
            let sz = group.len() as i64;
            for (i, &idx) in group.iter().enumerate() {
                res[idx] = total - prefix_total * 2 + idx as i64 * (2 * i as i64 - sz);
                prefix_total += idx as i64;
            }
        }
        res
    }
}