Skip to content

2121. Intervals Between Identical Elements 👍

  • 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
27
28
29
30
31
32
class Solution {
 public:
  vector<long long> getDistances(vector<int>& arr) {
    const int n = arr.size();
    vector<long long> ans(n);
    vector<long> prefix(n);
    vector<long> suffix(n);
    unordered_map<int, vector<int>> numToIndices;

    for (int i = 0; i < n; ++i)
      numToIndices[arr[i]].push_back(i);

    for (const auto& [_, indices] : numToIndices) {
      for (int i = 1; i < indices.size(); ++i) {
        const int currIndex = indices[i];
        const int prevIndex = indices[i - 1];
        prefix[currIndex] += prefix[prevIndex] + i * (currIndex - prevIndex);
      }
      for (int i = indices.size() - 2; i >= 0; --i) {
        const int currIndex = indices[i];
        const int prevIndex = indices[i + 1];
        suffix[currIndex] += suffix[prevIndex] +
                             (indices.size() - i - 1) * (prevIndex - currIndex);
      }
    }

    for (int i = 0; i < n; ++i)
      ans[i] = prefix[i] + suffix[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
29
30
31
32
class Solution {
  public long[] getDistances(int[] arr) {
    final int n = arr.length;
    long[] ans = new long[n];
    long[] prefix = new long[n];
    long[] suffix = new long[n];
    Map<Integer, List<Integer>> numToIndices = new HashMap<>();

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

    for (List<Integer> indices : numToIndices.values()) {
      for (int i = 1; i < indices.size(); ++i) {
        final int currIndex = indices.get(i);
        final int prevIndex = indices.get(i - 1);
        prefix[currIndex] += prefix[prevIndex] + i * (currIndex - prevIndex);
      }
      for (int i = indices.size() - 2; i >= 0; --i) {
        final int currIndex = indices.get(i);
        final int prevIndex = indices.get(i + 1);
        suffix[currIndex] += suffix[prevIndex] + (indices.size() - i - 1) * (prevIndex - currIndex);
      }
    }

    for (int i = 0; i < n; ++i)
      ans[i] = prefix[i] + suffix[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 getDistances(self, arr: list[int]) -> list[int]:
    prefix = [0] * len(arr)
    suffix = [0] * len(arr)
    numToIndices = collections.defaultdict(list)

    for i, a in enumerate(arr):
      numToIndices[a].append(i)

    for indices in numToIndices.values():
      for i in range(1, len(indices)):
        currIndex = indices[i]
        prevIndex = indices[i - 1]
        prefix[currIndex] += prefix[prevIndex] + i * (currIndex - prevIndex)
      for i in range(len(indices) - 2, -1, -1):
        currIndex = indices[i]
        prevIndex = indices[i + 1]
        suffix[currIndex] += (suffix[prevIndex] +
                              (len(indices) - i - 1) * (prevIndex - currIndex))

    return [p + s for p, s in zip(prefix, suffix)]