Java solution, O(n log n) using Map of Priority Queues

This actually solves the problem, including cases like [1 2 1 2 2] that most of the others here fail on, without going to an O(N^2) brute-force search.

public class Solution {
    public ArrayList<Integer> solve(ArrayList<Integer> A) {
        ArrayList<Integer> retval = new ArrayList<>();
        Map<Integer, PriorityQueue<Integer>> values = new HashMap<>();
        
        int index = 0;
        for (Integer curr : A) {
            retval.add(curr);
            if (!values.containsKey(curr)) {
                PriorityQueue<Integer> newq = new PriorityQueue<Integer>();
                newq.add(index);
                values.put(curr, newq);
            } else {
                PriorityQueue<Integer> currQ = values.get(curr);
                int indexToUpdate = currQ.poll();
                currQ.add(index);
                values.put(curr, currQ);
                PriorityQueue<Integer> updateQueue = values.getOrDefault(curr + 1, new PriorityQueue<>());
                updateQueue.add(indexToUpdate);
                values.put(curr + 1, updateQueue);
                retval.set(indexToUpdate, curr + 1);
            }
            ++index;
        }

        return retval;
    }
}

Yes, even the expected output for A=[1,2,1,2,2] is wrong. The expected output it is showing is [4,2,1,2,2] but should rather be [3,3,1,2,2].

Yes, I agree. I also implemented Heaps in Python and this solution is better. The IB’s Expected output is different for “Expected Output” Sections written in question part vs “Expected output” when you run the code with IB’s own solutions. And each of their solution (Fast, lightweight and Editorial) gives different output.
This question could be a good question of heaps and maps together but IB put it badly here.

Click here to start solving coding interview questions