Intervals Between Identical Elements LeetCode solution in C++

1
class Solution {
2
 public:
3
  vector<long long> getDistances(vector<int> &arr) {
4
    int size = arr.size();
5
    vector<long long> pre(size, 0), suf(size, 0), ans(size, 0);
6
    unordered_map<int, vector<int>> mp;
7

8
    // store index of the same elements.
9
    for (int i = 0; i < size; i++) mp[arr[i]].push_back(i);
10
    //  prefix will store the sum of absolute diff. from the starting index of
11
    //  current element to the current index.
12
    // pre[cur index]=pre[prev index]+ number of same elements before cur
13
    // element * abs diff of cur element pos and pre element pos
14
    for (auto &p : mp) {
15
      auto vec = p.second;
16
      for (int i = 1; i < vec.size(); i++)
17
        pre[vec[i]] = pre[vec[i - 1]] + i * (long)(vec[i] - vec[i - 1]);
18
    }
19
    // suffix will store the sum of absolute diff. from the ending index of
20
    // current element to the current index. same formula as above but in
21
    // reverse order
22
    for (auto &p : mp) {
23
      auto vec = p.second;
24
      for (int i = vec.size() - 2; i >= 0; i--)
25
        suf[vec[i]] = suf[vec[i + 1]] + long(vec.size() - 1 - i) * (long)(vec[i + 1] - vec[i]);
26
    }
27
    for (int i = 0; i < size; i++) ans[i] += (pre[i] + suf[i]);
28
    return ans;
29
  }
30
};