Leetcode #1443: Minimum Time to Collect All Apples in a Tree (https://interviewrecipes.com/leetcode-1443)

Leetcode_1443.cpp

class Solution {
    unordered_map<int, vector<int>> tree;
    unordered_map<int, bool> visited;
    vector<bool> hasApple;
    
    /** Convert edge list into adjacency list */
    void createTreeStructure(vector<vector<int>>& edges) {
        // Adjacency list representation.
        for (auto e: edges) {
            tree[e[0]].push_back(e[1]);   // a --> b
            tree[e[1]].push_back(e[0]);   // b --> a
        }
    }

    /** Returns time required to collect apples in the subtree of the 'node'. */
    int traverse(int node, int timeToVisitNode) {
        if (visited[node]) {          // Already visited this node and its subtree.
            return 0;                 // So no need to visit again.
        }
        visited[node] = true;         // Mark visited to avoid repetition.

        int subtreeTime = 0;          // Time required to collect apples in subtree.
        for (auto childNode: tree[node]) {
            subtreeTime += traverse(childNode, /* timeToVisitNode= */ 2);
        }

        // If the subtree has no apples, subtree time will be 0. And
        // if the current node also does not have an apple, then we can
        // simply avoid traversing the current branch of the tree to save
        // the time. In that case, time required for this subtree will be 0.
        if (subtreeTime == 0 && hasApple[node] == false) {
            return 0;
        }

        return subtreeTime + timeToVisitNode;
    }
public:

    int minTime(int n, vector<vector<int>>& edges, vector<bool>& hasApple) {
        this->hasApple = hasApple;
        createTreeStructure(edges);
        // Time to visit any node is 2 except 0.
        return traverse(/* startNode= */ 0, /* timeToVisitNode= */ 0);
    }
};