2121. 相同元素的间隔之和

140 阅读1分钟

题目:
给你一个下标从 0 开始、由 n 个整数组成的数组 arr 。

arr 中两个元素的 间隔 定义为它们下标之间的 绝对差 。更正式地,arr[i] 和 arr[j] 之间的间隔是 |i - j| 。

返回一个长度为 n 的数组 intervals ,其中 intervals[i] 是 **arr[i] **和 **arr **中每个相同元素(与 arr[i] 的值相同)的 间隔之和 

注意: |x| 是 x 的绝对值。
算法:
方法一:hash table
不出所料,TLE,尤其是当arr中的元素都相同时,时间复杂度N^2

func getDistances(arr []int) []int64 {
    ans := make([]int64, len(arr))
    m := make(map[int][]int, 0)
    for i := range arr {
        m[arr[i]] = append(m[arr[i]], i)
    }
    for i := range arr {
        if len(m[arr[i]]) == 1 {
            ans[i] = 0
            continue
        }
        sum := 0
        for _, j := range m[arr[i]] {
            sum = sum + abs(i - j)
        }
        ans[i] = int64(sum)
    }
    return ans
}

func abs(a int) int {
    if a < 0 {
        return -a
    }
    return a
}

方法二:优化
先将arr中相同元素的index放到list中,然后先求出list[0]的间隔和sum,移动到list第1个元素后,通过移动的间隔d = abs(list[1] - list[0])计算新的间隔和sum = sum + d * 左侧间隔和变化元素个数 - d * 右侧间隔和变化元素个数

不失一般性: d := abs(list[i] - list[i - 1]) 从i - 1 移动到i位置,
左侧间隔和增加:d * i
右侧间隔和减少:d * (n - i)
sum = sum + d * i - d * (n - i)

func getDistances(arr []int) []int64 {
    ans := make([]int64, len(arr))
    m := make(map[int][]int, 0)
    for i := range arr {
        m[arr[i]] = append(m[arr[i]], i)
    }
    for _, list := range m {
            sum := int64(0)
            for _, v := range list {
                sum = sum + int64(abs(v - list[0]))
            }
            ans[list[0]] = sum
            n := len(list)
            for j := 1; j < n; j ++ {
                d := abs(list[j] - list[j - 1])
                sum = sum + int64(j * d - (n - j) * d)
                ans[list[j]] = sum
            }
    }
    return ans
}

func abs(a int) int {
    if a < 0 {
        return -a
    }
    return a
}

方法三:前缀和
做前缀和re1[i]表示i之前的相同数到i位置的间隔和,re2[i]表示i之后相同数到i位置的间隔和,因此有: ans[i] = re1[i] + re2[i]

func getDistances(arr []int) []int64 {
    n := len(arr)
    ans := make([]int64, n)
    // m的value每保存上一次m的key出现位置和次数
    m := make(map[int][]int, 0)
    re1 := make([]int, n)
    re2 := make([]int, n)
    for i := 0; i < n; i ++ {
        if m[arr[i]] == nil {
            m[arr[i]] = make([]int, 2)
        }
        if i > 0 {
            // 左前缀和= 上次该数字出现位置的前缀和 + 次数 * 差
            re1[i] = re1[m[arr[i]][0]] + (i - m[arr[i]][0]) * m[arr[i]][1]
        }
        // 计算前缀和,更新上次出现位置,出现次数 + 1
        m[arr[i]][0] = i
        m[arr[i]][1] ++
    }
    m = make(map[int][]int, 0)
    for i := n - 1; i >= 0; i -- {
        if m[arr[i]] == nil {
            m[arr[i]] = make([]int, 2)
        }
        if i < n - 1 {
            // 左前缀和= 上次该数字出现位置的前缀和 + 次数 * 差
            re2[i] = re2[m[arr[i]][0]] + (m[arr[i]][0] - i) * m[arr[i]][1]
        }
        // 计算前缀和,更新上次出现位置,出现次数 + 1
        m[arr[i]][0] = i
        m[arr[i]][1] ++
    }
    for i := 0; i < n; i ++ {
        ans[i] = int64(re1[i] + re2[i])
    }
    return ans
}