leetcode-solutions 0.1.1

A collection of leetcode solution in rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

/*
 * No: 2121
 * Title: Intervals Between Identical Elements
 */

use crate::Solution;

impl Solution {
    pub fn get_distances(arr: Vec<i32>) -> Vec<i64> {
        use std::collections::HashMap;

        // For each element, record all the indexes where the same element
        // occurs
        let mut element_indexes: HashMap<i32, Vec<usize>> = HashMap::new();

        arr.iter().enumerate().for_each(|(idx, val)| {
            element_indexes.entry(*val).or_insert(Vec::new()).push(idx);
        });

        let mut res: Vec<i64> = vec![0; arr.len()];
        element_indexes.values().for_each(|indexes| {
            // prefix_sum: sum of all index diff from the starting point to the current index
            // suffix_sum: sum of all index diff from the current index to the end of the vec
            let mut prefix_sum: i64 = 0;
            let mut suffix_sum = indexes.iter().sum::<usize>() as i64;

            indexes.iter().enumerate().for_each(|(pos, index)| {
                suffix_sum -= *index as i64;
                res[*index] += suffix_sum - *index as i64 * (indexes.len() as i64 - pos as i64 - 1);
                res[*index] -= prefix_sum - *index as i64 * pos as i64;
                prefix_sum += *index as i64;
            });
        });

        res
    }
}