Leetcode 2368 - Reachable Nodes With Restrictions

Link to problem

Problem

There is an undirected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given a 2D integer array edges of length n - 1 where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree. You are also given an integer array restricted which represents restricted nodes.

Return the maximum number of nodes you can reach from node 0 without visiting a restricted node.

Note that node 0 will not be a restricted node.

Solution

A DFS will do :).

First use a hash map to store the tree using the edges provided for later traversal. Note that since the tree is undirected, we need to store both direction of an edge. Moreover, we will need a hash map for restricted nodes and another to keep track of visited nodes (since the tree is undirected, we don’t want to visit a node multiple times).

Implementation

class Solution {
public:
    int reachableNodes(
            int n,
            vector<vector<int>>& edges,
            vector<int>& restricted)
    {
        unordered_set<int> restrict_set(restricted.begin(), restricted.end());
        unordered_map<int, vector<int>> tree;
        unordered_set<int> visited;
        int output = 1;

        for (const auto &edge: edges) {
            tree[edge[0]].push_back(edge[1]);
            tree[edge[1]].push_back(edge[0]);
        }

        stack<int> temp;
        temp.push(0);
        while (!temp.empty()) {
            int cur = temp.top();
            temp.pop();
            visited.insert(cur);

            for (auto child: tree[cur]) {
                if (restrict_set.find(child) == restrict_set.end() &&
                    visited.find(child) == visited.end()) {
                    temp.push(child);
                    output++;
                }
            }
        }

        return output;
    }
};