Filter Restaurants by Vegan-Friendly, Price and Distance

Problem

Given the array restaurants where restaurants[i] = [idi, ratingi, veganFriendlyi, pricei, distancei]. You have to filter the restaurants using three filters.

The veganFriendly filter will be either true (meaning you should only include restaurants with veganFriendlyi set to true) or false (meaning you can include any restaurant). In addition, you have the filters maxPrice and maxDistance which are the maximum value for price and distance of restaurants you should consider respectively.

Return the array of restaurant IDs after filtering, ordered by rating from highest to lowest. For restaurants with the same rating, order them by id from highest to lowest. For simplicity veganFriendlyi and veganFriendly take value 1 when it is true , and 0 when it is false.

Examples

Example 1


    
    
    Input: restaurants = [[1,4,1,40,10],[2,8,0,50,5],[3,8,1,30,4],[4,10,0,10,3],[5,1,1,15,1]], veganFriendly = 1, maxPrice = 50, maxDistance = 10
    Output: [3,1,5] 
    Explanation: The restaurants are:
    Restaurant 1 [id=1, rating=4, veganFriendly=1, price=40, distance=10]
    Restaurant 2 [id=2, rating=8, veganFriendly=0, price=50, distance=5]
    Restaurant 3 [id=3, rating=8, veganFriendly=1, price=30, distance=4]
    Restaurant 4 [id=4, rating=10, veganFriendly=0, price=10, distance=3]
    Restaurant 5 [id=5, rating=1, veganFriendly=1, price=15, distance=1] 
    After filter restaurants with veganFriendly = 1, maxPrice = 50 and maxDistance = 10 we have restaurant 3, restaurant 1 and restaurant 5 (ordered by rating from highest to lowest).

Example 2


    
    
    Input: restaurants = [[1,4,1,40,10],[2,8,0,50,5],[3,8,1,30,4],[4,10,0,10,3],[5,1,1,15,1]], veganFriendly = 0, maxPrice = 50, maxDistance = 10
    Output: [4,3,2,1,5]
    Explanation: The restaurants are the same as in example 1, but in this case the filter veganFriendly = 0, therefore all restaurants are considered.

Example 3


    
    
    Input: restaurants = [[1,4,1,40,10],[2,8,0,50,5],[3,8,1,30,4],[4,10,0,10,3],[5,1,1,15,1]], veganFriendly = 0, maxPrice = 30, maxDistance = 3
    Output: [4,5]

Constraints

1 <= restaurants.length <= 10^4
restaurants[i].length == 5
1 <= idi, ratingi, pricei, distancei <= 10^5
1 <= maxPrice, maxDistance <= 10^5
veganFriendlyi and veganFriendly are 0 or 1.
All idi are distinct.

Solution

Method 1 – Filtering and Sorting

Intuition

We filter the restaurants based on the vegan, price, and distance constraints, then sort the filtered list by rating (descending) and id (descending).

Approach

Filter the restaurants:
- If veganFriendly == 1, only include those with veganFriendlyi == 1.
- Only include those with pricei <= maxPrice and distancei <= maxDistance.
Sort the filtered restaurants by rating (descending), then by id (descending).
Return the list of ids in the sorted order.

Code

C++

class Solution {
public:
    vector<int> filterRestaurants(vector<vector<int>>& restaurants, int veganFriendly, int maxPrice, int maxDistance) {
        vector<vector<int>> filtered;
        for (auto& r : restaurants) {
            if ((veganFriendly == 0 || r[2] == 1) && r[3] <= maxPrice && r[4] <= maxDistance) {
                filtered.push_back(r);
            }
        }
        sort(filtered.begin(), filtered.end(), [](const vector<int>& a, const vector<int>& b) {
            if (a[1] != b[1]) return a[1] > b[1];
            return a[0] > b[0];
        });
        vector<int> ans;
        for (auto& r : filtered) ans.push_back(r[0]);
        return ans;
    }
};

Go

func filterRestaurants(restaurants [][]int, veganFriendly, maxPrice, maxDistance int) []int {
    filtered := [][]int{}
    for _, r := range restaurants {
        if (veganFriendly == 0 || r[2] == 1) && r[3] <= maxPrice && r[4] <= maxDistance {
            filtered = append(filtered, r)
        }
    }
    sort.Slice(filtered, func(i, j int) bool {
        if filtered[i][1] != filtered[j][1] {
            return filtered[i][1] > filtered[j][1]
        }
        return filtered[i][0] > filtered[j][0]
    })
    ans := make([]int, len(filtered))
    for i, r := range filtered {
        ans[i] = r[0]
    }
    return ans
}

Java

class Solution {
    public List<Integer> filterRestaurants(int[][] restaurants, int veganFriendly, int maxPrice, int maxDistance) {
        List<int[]> filtered = new ArrayList<>();
        for (int[] r : restaurants) {
            if ((veganFriendly == 0 || r[2] == 1) && r[3] <= maxPrice && r[4] <= maxDistance) {
                filtered.add(r);
            }
        }
        filtered.sort((a, b) -> a[1] != b[1] ? b[1] - a[1] : b[0] - a[0]);
        List<Integer> ans = new ArrayList<>();
        for (int[] r : filtered) ans.add(r[0]);
        return ans;
    }
}

Kotlin

class Solution {
    fun filterRestaurants(restaurants: Array<IntArray>, veganFriendly: Int, maxPrice: Int, maxDistance: Int): List<Int> {
        return restaurants.filter {
            (veganFriendly == 0 || it[2] == 1) && it[3] <= maxPrice && it[4] <= maxDistance
        }.sortedWith(compareByDescending<IntArray> { it[1] }.thenByDescending { it[0] })
         .map { it[0] }
    }
}

Python

class Solution:
    def filterRestaurants(self, restaurants: list[list[int]], veganFriendly: int, maxPrice: int, maxDistance: int) -> list[int]:
        filtered = [r for r in restaurants if (veganFriendly == 0 or r[2] == 1) and r[3] <= maxPrice and r[4] <= maxDistance]
        filtered.sort(key=lambda x: (x[1], x[0]), reverse=True)
        return [r[0] for r in filtered]

Rust

impl Solution {
    pub fn filter_restaurants(restaurants: Vec<Vec<i32>>, vegan_friendly: i32, max_price: i32, max_distance: i32) -> Vec<i32> {
        let mut filtered: Vec<_> = restaurants.into_iter().filter(|r| {
            (vegan_friendly == 0 || r[2] == 1) && r[3] <= max_price && r[4] <= max_distance
        }).collect();
        filtered.sort_by(|a, b| b[1].cmp(&a[1]).then(b[0].cmp(&a[0])));
        filtered.into_iter().map(|r| r[0]).collect()
    }
}

TypeScript

class Solution {
    filterRestaurants(restaurants: number[][], veganFriendly: number, maxPrice: number, maxDistance: number): number[] {
        return restaurants.filter(r => (veganFriendly === 0 || r[2] === 1) && r[3] <= maxPrice && r[4] <= maxDistance)
            .sort((a, b) => b[1] - a[1] || b[0] - a[0])
            .map(r => r[0]);
    }
}

Complexity

⏰ Time complexity: O(n log n), where n is the number of restaurants, due to sorting.
🧺 Space complexity: O(n), for storing the filtered list.