https://leetcode.com/problems/campus-bikes

gistfile1.txt

class Solution {
    public int[] assignBikes(int[][] workers, int[][] bikes) {
        Node[] res = new Node[workers.length];
        Set<Integer> bikeSet = new HashSet<>();
        Set<Integer> workerSet = new HashSet<>();
        
        //create a heap as per condition described in the question.
        Queue<Node> pq = new PriorityQueue<>(new Comparator<Node>(){
           
            public int compare(Node n1, Node n2) {
                
                
                int d1 = getDistance(n1.worker, n1.bike);
                int d2 = getDistance(n2.worker, n2.bike);
                if(d1==d2) {
                    if(n1.workerIndex == n2.workerIndex) {
                        return n1.bikeIndex - n2.bikeIndex;    
                    } return n1.workerIndex - n2.workerIndex;
                    
                } else return d1-d2;
            }
        });
        
        //add cartesian product of workers and bikes to the heap.
        for(int i = 0;i<workers.length;i++) {
            for(int j = 0;j<bikes.length;j++) {
                pq.add(new Node(workers[i], bikes[j], i, j));
            }
        }
        int[] result = new int[workers.length];
        // take each node off the heap and only insert it as part of result only if the worker and bike have not been already used
        int index = 0;
        while(!pq.isEmpty()) {
            Node n = pq.poll();
            if(!workerSet.contains(n.workerIndex) && !bikeSet.contains(n.bikeIndex)) {
                result[n.workerIndex] = n.bikeIndex;
                workerSet.add(n.workerIndex);
                bikeSet.add(n.bikeIndex);
            }
        }
        return result;
        
    }
    
    private int getDistance(int[] x, int[] y) {
        return Math.abs(x[0]-y[0]) + Math.abs(x[1]-y[1]);
    }
    
    class Node {
        int[] worker;
        int[] bike;
        int workerIndex;
        int bikeIndex;
        
        Node(int[] worker, int[] bike, int workerIndex, int bikeIndex) {
            this.worker = worker;
            this.bike = bike;
            this.workerIndex = workerIndex;
            this.bikeIndex = bikeIndex;
        }
    }
}