Description
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 <= 1050 <= nums[i] <= 109
Note: This question is the same as 2121: Intervals Between Identical Elements.
Solutions
Solution 1: Hash Map + Prefix Sum
First, use a hash map d to record the list of indices for each element in the array nums, that is, d[x] represents the list of all indices in nums where the value is x.
For each list of indices idx in the hash map d, we can calculate the value of arr[i] for each index i in idx. For the first index idx[0], the sum of distances to all indices on the right is right = ∑_{i=0}m-1 - idx[0] × m. Then, we iterate through idx, and for each iteration, compute ans[idx[i]] = left + right, then update left and right as follows: left = left + (idx[i+1] - idx[i]) × (i+1), and right = right - (idx[i+1] - idx[i]) × (m-i-1).
After the iteration, we obtain the array arr corresponding to each element in nums, which is stored in ans.
The time complexity is O(n) and the space complexity is O(n), where n is the length of the array nums.
PythonJavaC++GoTypeScriptRust
class Solution:
def distance(self, nums: List[int]) -> List[int]:
d = defaultdict(list)
for i, x in enumerate(nums):
d[x].append(i)
ans = [0] * len(nums)
for idx in d.values():
left, right = 0, sum(idx) - len(idx) * idx[0]
for i in range(len(idx)):
ans[idx[i]] = left + right
if i + 1 < len(idx):
left += (idx[i + 1] - idx[i]) * (i + 1)
right -= (idx[i + 1] - idx[i]) * (len(idx) - i - 1)
return ans(code-box)
class Solution {
public long[] distance(int[] nums) {
int n = nums.length;
long[] ans = new long[n];
Map<Integer, List<Integer>> d = new HashMap<>();
for (int i = 0; i < n; ++i) {
d.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
}
for (var idx : d.values()) {
int m = idx.size();
long left = 0;
long right = -1L * m * idx.get(0);
for (int i : idx) {
right += i;
}
for (int i = 0; i < m; ++i) {
ans[idx.get(i)] = left + right;
if (i + 1 < m) {
left += (idx.get(i + 1) - idx.get(i)) * (i + 1L);
right -= (idx.get(i + 1) - idx.get(i)) * (m - i - 1L);
}
}
}
return ans;
}
}(code-box)
class Solution {
public:
vector<long long> distance(vector<int>& nums) {
int n = nums.size();
vector<long long> ans(n);
unordered_map<int, vector<int>> d;
for (int i = 0; i < n; ++i) {
d[nums[i]].push_back(i);
}
for (auto& [_, idx] : d) {
int m = idx.size();
long long left = 0;
long long right = -1LL * m * idx[0];
for (int i : idx) {
right += i;
}
for (int i = 0; i < m; ++i) {
ans[idx[i]] = left + right;
if (i + 1 < m) {
left += (idx[i + 1] - idx[i]) * (i + 1);
right -= (idx[i + 1] - idx[i]) * (m - i - 1);
}
}
}
return ans;
}
};(code-box)
func distance(nums []int) []int64 {
n := len(nums)
ans := make([]int64, n)
d := map[int][]int{}
for i, x := range nums {
d[x] = append(d[x], i)
}
for _, idx := range d {
m := len(idx)
left, right := 0, -m*idx[0]
for _, i := range idx {
right += i
}
for i := range idx {
ans[idx[i]] = int64(left + right)
if i+1 < m {
left += (idx[i+1] - idx[i]) * (i + 1)
right -= (idx[i+1] - idx[i]) * (m - i - 1)
}
}
}
return ans
}(code-box)
function distance(nums: number[]): number[] {
const n = nums.length;
const ans = new Array(n).fill(0);
const d = new Map<number, number[]>();
for (let i = 0; i < n; ++i) {
if (!d.has(nums[i])) {
d.set(nums[i], []);
}
d.get(nums[i])!.push(i);
}
for (const idx of d.values()) {
const m = idx.length;
let left = 0;
let right = -1 * m * idx[0];
for (const i of idx) {
right += i;
}
for (let i = 0; i < m; ++i) {
ans[idx[i]] = left + right;
if (i + 1 < m) {
left += (idx[i + 1] - idx[i]) * (i + 1);
right -= (idx[i + 1] - idx[i]) * (m - i - 1);
}
}
}
return ans;
}(code-box)
use std::collections::HashMap;
impl Solution {
pub fn distance(nums: Vec<i32>) -> Vec<i64> {
let n = nums.len();
let mut ans = vec![0i64; n];
let mut d: HashMap<i32, Vec<usize>> = HashMap::new();
for i in 0..n {
d.entry(nums[i]).or_insert(Vec::new()).push(i);
}
for idx in d.values() {
let m = idx.len();
let mut left = 0i64;
let mut right = 0i64;
for &i in idx {
right += i as i64;
}
right -= m as i64 * idx[0] as i64;
for i in 0..m {
ans[idx[i]] = left + right;
if i + 1 < m {
let diff = (idx[i + 1] - idx[i]) as i64;
left += diff * (i + 1) as i64;
right -= diff * (m - i - 1) as i64;
}
}
}
ans
}
}(code-box)
