Skip to content

2615. Sum of Distances 👍

  • Time: $O(n)$
  • Space: $O(n)$
 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
class Solution {
 public:
  vector<long long> distance(vector<int>& nums) {
    vector<long long> ans(nums.size());
    unordered_map<int, vector<int>> numToIndices;

    for (int i = 0; i < nums.size(); ++i)
      numToIndices[nums[i]].push_back(i);

    for (const auto& [_, indices] : numToIndices) {
      const int n = indices.size();
      if (n == 1)
        continue;
      long long sumSoFar = accumulate(indices.begin(), indices.end(), 0LL);
      int prevIndex = 0;
      for (int i = 0; i < n; ++i) {
        sumSoFar += (i - 1) * (indices[i] - prevIndex);
        sumSoFar -= (n - 1 - i) * (indices[i] - prevIndex);
        ans[indices[i]] = sumSoFar;
        prevIndex = indices[i];
      }
    }

    return ans;
  }
};

 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
public class Solution {
  public long[] distance(int[] nums) {
    long[] ans = new long[nums.length];
    Map<Integer, List<Integer>> numToIndices = new HashMap<>();

    for (int i = 0; i < nums.length; ++i) {
      numToIndices.putIfAbsent(nums[i], new ArrayList<>());
      numToIndices.get(nums[i]).add(i);
    }

    for (List<Integer> indices : numToIndices.values()) {
      final int n = indices.size();
      if (n == 1) {
        continue;
      }
      long sumSoFar = indices.stream().mapToLong(Integer::intValue).sum();
      int prevIndex = 0;
      for (int i = 0; i < n; ++i) {
        sumSoFar += (i - 1) * (indices.get(i) - prevIndex);
        sumSoFar -= (n - 1 - i) * (indices.get(i) - prevIndex);
        ans[indices.get(i)] = sumSoFar;
        prevIndex = indices.get(i);
      }
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
  def distance(self, nums: List[int]) -> List[int]:
    ans = [0] * len(nums)
    numToIndices = collections.defaultdict(list)

    for i, num in enumerate(nums):
      numToIndices[num].append(i)

    for indices in numToIndices.values():
      n = len(indices)
      if n == 1:
        continue
      sumSoFar = sum(indices)
      prevIndex = 0
      for i in range(n):
        sumSoFar += (i - 1) * (indices[i] - prevIndex)
        sumSoFar -= (n - 1 - i) * (indices[i] - prevIndex)
        ans[indices[i]] = sumSoFar
        prevIndex = indices[i]

    return ans