Filter Restaurants by Vegan-Friendly, Price and Distance LeetCode solution in Java

1
class Restaurant {
2
  int id, rating;
3

4
  Restaurant(int id, int rating) {
5
    this.id = id;
6
    this.rating = rating;
7
  }
8
}
9

10
class RestaurantComparator implements Comparator<Restaurant> {
11
  @Override
12
  public int compare(Restaurant r1, Restaurant r2) {
13
    return r1.rating == r2.rating ? r2.id - r1.id : r2.rating - r1.rating;
14
  }
15
}
16

17
class Solution {
18
  public List<Integer> filterRestaurants(
19
      int[][] restaurants, int veganFriendly, int maxPrice, int maxDistance) {
20
    PriorityQueue<Restaurant> heap = new PriorityQueue<>(new RestaurantComparator());
21
    if (veganFriendly == 1) {
22
      for (int[] restaurant : restaurants) {
23
        if (restaurant[2] == 1 && restaurant[3] <= maxPrice && restaurant[4] <= maxDistance) {
24
          heap.offer(new Restaurant(restaurant[0], restaurant[1]));
25
        }
26
      }
27
    } else {
28
      for (int[] restaurant : restaurants) {
29
        if (restaurant[3] <= maxPrice && restaurant[4] <= maxDistance) {
30
          heap.offer(new Restaurant(restaurant[0], restaurant[1]));
31
        }
32
      }
33
    }
34
    List<Integer> answer = new ArrayList<>();
35
    while (!heap.isEmpty()) {
36
      answer.add(heap.poll().id);
37
    }
38
    return answer;
39
  }
40
}