[LC] 1057. Campus Bikes

Posted 2021-03-06 xuanlu

On a campus represented as a 2D grid, there are N workers and Mbikes, with N <= M. Each worker and bike is a 2D coordinate on this grid.

Our goal is to assign a bike to each worker. Among the available bikes and workers, we choose the (worker, bike) pair with the shortest Manhattan distance between each other, and assign the bike to that worker. (If there are multiple (worker, bike) pairs with the same shortest Manhattan distance, we choose the pair with the smallest worker index; if there are multiple ways to do that, we choose the pair with the smallest bike index). We repeat this process until there are no available workers.

The Manhattan distance between two points p1 and p2 is Manhattan(p1, p2) = |p1.x - p2.x| + |p1.y - p2.y|.

Return a vector ans of length N, where ans[i] is the index (0-indexed) of the bike that the i-th worker is assigned to.

 

class Solution {
    public int[] assignBikes(int[][] workers, int[][] bikes) {
        int[] res = new int[workers.length];
        Arrays.fill(res, -1);
        boolean[] usedBikes = new boolean[bikes.length];
        PriorityQueue<Cell> pq = new PriorityQueue<>(new Comparator<>() {
            @Override
            public int compare(Cell c1, Cell c2) {
                if (c1.dist == c2.dist) {
                    if (c1.worker == c2.worker) {
                        return c1.bike - c2.bike;
                    }
                    return c1.worker - c2.worker;
                }
                return c1.dist - c2.dist;
            }
        });
        
        for (int i = 0; i < workers.length; i++) {
            for (int j = 0; j < bikes.length; j++) {
                pq.offer(new Cell(i, j, getDist(workers[i], bikes[j])));
            }
        }
        
        int count = 0;
        while (count < workers.length) {
            Cell cur = pq.poll();
            int workerIndex = cur.worker;
            int bikeIndex = cur.bike;
            if (!usedBikes[bikeIndex] && res[workerIndex] == -1) {
                res[workerIndex] = bikeIndex;
                count += 1;
                usedBikes[bikeIndex] = true;
            }
        }
        return res;
    }
    
    private int getDist(int[] worker, int[] bike) {
        return Math.abs(worker[0] - bike[0]) + Math.abs(worker[1] - bike[1]);
    }
}

class Cell {
    int worker;
    int bike;
    int dist;
    public Cell(int worker, int bike, int dist) {
        this.worker = worker;
        this.bike = bike;
        this.dist = dist;
    }
}