Intervals Between Identical Elements

medium array hash table prefix sum

Problem

Given a 0-indexed array arr of n integers, return an array intervals of length n where intervals[i] is the sum of |i − j| over every index j such that arr[j] == arr[i] (including only matching values, and excluding i itself since |i − i| = 0).

Inputarr = [2,1,3,1,2,3,3]

Output[4,2,7,2,4,4,5]

Index 2 holds value 3, which also sits at indices 5 and 6: |2−5| + |2−6| = 3 + 4 = 7.

arr (comma-separated)

def getDistances(arr):
    n = len(arr)
    res = [0] * n
    groups = {}                       # value -> list of indices
    for i, v in enumerate(arr):
        groups.setdefault(v, []).append(i)
    for idxs in groups.values():
        m = len(idxs)
        left = 0                      # sum of indices seen so far
        right = sum(idxs)             # sum of all indices in group
        for k in range(m):
            idx = idxs[k]
            right -= idx              # indices strictly after k
            # k items before, m-1-k items after
            res[idx] = (k * idx - left) + (right - (m - 1 - k) * idx)
            left += idx               # idx now counts as "before"
    return res

function getDistances(arr) {
  const n = arr.length;
  const res = new Array(n).fill(0);
  const groups = new Map();           // value -> array of indices
  for (let i = 0; i < n; i++) {
    if (!groups.has(arr[i])) groups.set(arr[i], []);
    groups.get(arr[i]).push(i);
  }
  for (const idxs of groups.values()) {
    const m = idxs.length;
    let left = 0;                      // sum of indices seen so far
    let right = idxs.reduce((a, b) => a + b, 0);
    for (let k = 0; k < m; k++) {
      const idx = idxs[k];
      right -= idx;                    // indices strictly after k
      res[idx] = (k * idx - left) + (right - (m - 1 - k) * idx);
      left += idx;                     // idx now counts as "before"
    }
  }
  return res;
}

long[] getDistances(int[] arr) {
    int n = arr.length;
    long[] res = new long[n];
    Map<Integer, List<Integer>> groups = new HashMap<>();
    for (int i = 0; i < n; i++)
        groups.computeIfAbsent(arr[i], z -> new ArrayList<>()).add(i);
    for (List<Integer> idxs : groups.values()) {
        int m = idxs.size();
        long left = 0, right = 0;       // running prefix / suffix sums
        for (int x : idxs) right += x;
        for (int k = 0; k < m; k++) {
            long idx = idxs.get(k);
            right -= idx;               // indices strictly after k
            res[(int) idx] = (k * idx - left) + (right - (long)(m - 1 - k) * idx);
            left += idx;                // idx now counts as "before"
        }
    }
    return res;
}

vector<long long> getDistances(vector<int>& arr) {
    int n = arr.size();
    vector<long long> res(n, 0);
    unordered_map<int, vector<int>> groups;
    for (int i = 0; i < n; i++) groups[arr[i]].push_back(i);
    for (auto& kv : groups) {
        auto& idxs = kv.second;
        int m = idxs.size();
        long long left = 0, right = 0;  // running prefix / suffix sums
        for (int x : idxs) right += x;
        for (int k = 0; k < m; k++) {
            long long idx = idxs[k];
            right -= idx;               // indices strictly after k
            res[idx] = (k * idx - left) + (right - (long long)(m - 1 - k) * idx);
            left += idx;                // idx now counts as "before"
        }
    }
    return res;
}

Explanation

The naive idea — for each i, scan the whole array and add |i − j| whenever arr[j] == arr[i] — is O(n²) and too slow for n up to 10⁵. The fix is to notice that only elements of the same value ever matter, so we handle each value independently.

First we group indices by value into a hash map: groups[v] is the list of all positions holding value v. Because we read arr left to right, every such list comes out already sorted in increasing order.

Within one group of sorted indices idxs of size m, fix the element at position k with index value idx. Every group member sits either before it or after it. For the k members before it, each distance is idx − idxs[j], and the total is k·idx − left, where left is the running sum of the indices already passed. For the m−1−k members after it, each distance is idxs[j] − idx, totalling right − (m−1−k)·idx, where right is the sum of the indices still ahead.

So with two running prefix sums — left (sum of indices already visited) and right (sum of indices not yet visited) — each element's answer is computed in O(1). We slide idx from the right bucket into the left bucket as we advance.

Example: for value 3 in [2,1,3,1,2,3,3] the indices are [2,5,6]. At k=0, idx=2: nothing before, after-sum is 5+6=11 with two items, so res[2] = 11 − 2·2 = 7. That matches the expected output.

Time: O(n) Space: O(n)