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

1
// comparator class for sorting restaurants by their rating
2
class comp {
3
 public:
4
  bool operator()(vector<int> a, vector<int> b) {
5
    // same rating then sort by ids
6
    if (a[1] == b[1]) return a[0] > b[0];
7
    return a[1] > b[1];
8
  }
9
};
10
class Solution {
11
 public:
12
  vector<int> filterRestaurants(vector<vector<int>> &restaurants, int veganFriendly, int maxPrice,
13
                                int maxDistance) { // sort restaurants by their rating
14
    sort(restaurants.begin(), restaurants.end(), comp());
15
    vector<int> ans;
16
    for (int i = 0; i < restaurants.size(); i++) {
17
      // veganfriendly
18
      if (veganFriendly) {
19
        // store ids who satisfy the constraints
20
        if (restaurants[i][2] == veganFriendly and restaurants[i][3] <= maxPrice and
21
            restaurants[i][4] <= maxDistance) {
22
          ans.push_back(restaurants[i][0]);
23
        }
24
      }
25
      // non vegan
26
      else {
27
        // store ids who satisfy the constraints
28
        if (restaurants[i][3] <= maxPrice and restaurants[i][4] <= maxDistance) {
29
          ans.push_back(restaurants[i][0]);
30
        }
31
      }
32
    }
33
    return ans;
34
  }
35
};