leetcode2121.java

/// https://leetcode.com/problems/intervals-between-identical-elements
///
class Solution {
    public long[] getDistances(int[] arr) {
        int n = arr.length;
        HashMap<Integer, ArrayList<Integer>> map = new HashMap<>();
        for (int i=0; i<n; i++) {
            ArrayList<Integer> a = map.get(arr[i]);
            if (a == null) {
                a = new ArrayList<>();
            }
            a.add(i);
            map.put(arr[i], a);
        }
        for (int num : map.keySet()) {
            ArrayList<Integer> a = map.get(num);
            int prev = a.get(0);
            for (int i=1; i<a.size(); i++) {
                a.set(i-1, a.get(i)-a.get(i-1));
            }
            a.remove(a.size()-1);
        }
        HashMap<Integer, ArrayList<Long>> left = new HashMap<>();
        HashMap<Integer, ArrayList<Long>> right = new HashMap<>();
        HashMap<Integer, Integer> ptr = new HashMap<>();
        for (int num : map.keySet()) {
            ArrayList<Integer> intervals = map.get(num);
            int ints_sz = intervals.size();
            ArrayList<Long> l = new ArrayList<>();
            int cnt = 1;
            for (int i=ints_sz-1; i>=0; i--) {
                l.add(0, ((long)intervals.get(i))*cnt + (l.size() > 0 ? l.get(0) : 0));
                cnt++;
            }
            left.put(num, l);
            cnt = 1;
            ArrayList<Long> r = new ArrayList<>();
            for (int i=0; i<ints_sz; i++) {
                r.add(r.size(), intervals.get(i)*cnt + (r.size() > 0 ? r.get(r.size()-1) : 0));
                cnt++;
            }
            right.put(num, r);
            ptr.put(num, 0);
        }
        long[] res = new long[n];
        for (int i = 0; i < n; i++) {
            int num = arr[i];
            ArrayList<Long> l = left.get(num);
            ArrayList<Long> r = right.get(num);
            int curr_ptr = ptr.get(num);
            res[i] = (curr_ptr < l.size() ? l.get(curr_ptr) : 0)
                + (curr_ptr-1>=0 ? r.get(curr_ptr-1) : 0);
            ptr.put(num, curr_ptr+1);
        }
        return res;
    }
}