2121. 相同元素的间隔之和

2121. 相同元素的间隔之和

1. 题意

给定一个数组 arr ，定义两个元素的间隔是数组中与该元素的值相同的位置之差。
即：dis = | i - j | ，(arr[i] == arr[j])
返回每位数的间隔之和

---

样例

输入：arr = [2,1,3,1,2,3,3]
输出：[4,2,7,2,4,4,5]
解释：
下标 0 ：另一个 2 在下标 4 ，|0 - 4| = 4
下标 1 ：另一个 1 在下标 3 ，|1 - 3| = 2
下标 2 ：另两个 3 在下标 5 和 6 ，|2 - 5| + |2 - 6| = 7
下标 3 ：另一个 1 在下标 1 ，|3 - 1| = 2
下标 4 ：另一个 2 在下标 0 ，|4 - 0| = 4
下标 5 ：另两个 3 在下标 2 和 6 ，|5 - 2| + |5 - 6| = 4
下标 6 ：另两个 3 在下标 2 和 5 ，|6 - 2| + |6 - 5| = 5

输入：arr = [10,5,10,10]
输出：[5,0,3,4]
解释：
下标 0 ：另两个 10 在下标 2 和 3 ，|0 - 2| + |0 - 3| = 5
下标 1 ：只有这一个 5 在数组中，所以到相同元素的间隔之和是 0
下标 2 ：另两个 10 在下标 0 和 3 ，|2 - 0| + |2 - 3| = 3
下标 3 ：另两个 10 在下标 0 和 2 ，|3 - 0| + |3 - 2| = 4

---

2. 思路

对每位数拉链，vector存储元素 x 的依次的位置

1. 前缀和 + 后缀和

/* 相同元素的下标为数组p: p0   p1   p2   p3   p4
   两两间隔为:               a    b   c    d
   ans[0] = p4-p0 + p3-p0 + p2-p0 + p1 - p0
          = 0           +    4a+3b+2c+1d  
   ans[1] = a           +       3b+2c+1d
   ans[2] = a+2b        +          2c+1d
   ans[3] = a+2b+3c     +             1d
   ans[4] = a+2b+3c+4d  +              0
   从上可以看出规律, 使用前缀和leftSum和后缀和rightSum, 进行优化计算;
*/

2. 一遍递推过去

• 先遍历一次 v, 计算出第一个元素（位置为 0）到其余元素的间隔和

• 第二次遍历，不断向右计算下一个元素的间隔和。

  • 第 i - 1 个元素到第 i 个元素，间隔和变化：

    • 左边有 i 个元素，间隔和变大 v[i] - v[i - 1]
    • 右边有 n - i 个元素， 间隔和变小 v[i] - v[i - 1]
    • 总间隔和变化量：（2 * i - n） * （v[i] - v[i - 1]）

---

3. 算法

哈希 + 前缀和

---

4. 代码

代码1

class Solution {
public:
    typedef long long ll;
    vector<long long> getDistances(vector<int>& arr) {
        vector<ll> ans(arr.size());
        map<int, vector<ll> > mp;
        for (int i = 0; i < arr.size(); i++) {
            int x = arr[i];
            mp[x].push_back(i);
        }
        for (auto &[k, v]: mp) {
            int n = v.size();
            ll suffix = 0, prefix = 0; // 后缀 & 前缀
            vector<ll> diff(n); // 相邻元素间隔差
            for (int i = 0; i < n - 1; i++)  
                diff[i] = v[i + 1] - v[i];
            for (int i = n - 2; i >= 0; i--) // 计算后缀 4a + 3b + 2c + d
                suffix += diff[i] * (n - i - 1);

            ans[v[0]] = suffix; // 第一个位置的值就是整个后缀
            for (int i = 1; i < n; i++) {
                prefix += i * diff[i - 1]; // 前缀计算
                suffix -= (n - i) * diff[i - 1]; // 后缀变化
                ans[v[i]] = prefix + suffix; // 每一位的值
            }
        }
        return ans;
    }
};

代码2

class Solution {
public:
    typedef long long ll;
    vector<long long> getDistances(vector<int>& arr) {
        vector<ll> ans(arr.size());
        map<int, vector<ll> > mp;
        for (int i = 0; i < arr.size(); i++) {
            mp[arr[i]].push_back(i);
        }
        for (auto &[k, v]: mp) {
            int n = v.size();
            ll sum = 0;
            for (int i = 0; i < n; i++)
                sum += v[i] - v[0];
            
            ans[v[0]] = sum;

            for (int i = 1; i < n; i++) {
                sum += (2 * i - n) * (v[i] - v[i - 1]);
                ans[v[i]] = sum;
            }
        }
        return ans;
    }
};