LeetCode 863: All Nodes Distance K in Binary Tree

Question

I'm posting two similar solutions for LeetCode's "All Nodes Distance K in Binary Tree". If you'd like to review, please do so. Thank you!

Problem

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.

Input:

• root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

Output:

• [7,4,1]

Explanation:

• The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

• Note that the inputs "root" and "target" are actually TreeNodes.
• The descriptions of the inputs above are just serializations of these objects.

Note:

• The given tree is non-empty.
• Each node in the tree has unique values 0 <= node.val <= 500.
• The target node is a node in the tree.
• 0 <= K <= 1000.

Note that the inputs "root" and "target" are actually TreeNodes. The descriptions of the inputs above are just serializations of these objects.

Solution 1

// The following block might slightly improve the execution time;
// Can be removed;
static const auto __optimize__ = []() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    return 0;
}();

// Most of headers are already included;
// Can be removed;
#include <cstdint>
#include <vector>
#include <unordered_map>
#include <unordered_set>

using ValueType = int;

static const struct Solution {
        static const std::vector<ValueType> distanceK(
            TreeNode* root,
            TreeNode* target,
            const ValueType K
        ) {
            std::vector<ValueType> res;
            std::unordered_map<TreeNode*, TreeNode*> parents;
            std::unordered_set<TreeNode*> visited;


            getParent(root, parents);
            depthFirstSearch(target, K, parents, visited, res);
            return res;
        }

    private:
        static const void getParent(
            TreeNode* node,
            std::unordered_map<TreeNode*, TreeNode*>& parents
        ) {
            if (!node) {
                return;
            }

            if (node->left) {
                parents[node->left] = node;
                getParent(node->left, parents);
            }

            if (node->right) {
                parents[node->right] = node;
                getParent(node->right, parents);
            }
        }


        static const void depthFirstSearch(
            TreeNode* node,
            const ValueType K,
            std::unordered_map<TreeNode*, TreeNode*>& parents,
            std::unordered_set<TreeNode*>& visited,
            std::vector<ValueType>& res
        ) {
            if (!node) {
                return;
            }

            if (visited.count(node) > 0) {
                return;
            }

            visited.insert(node);

            if (!K) {
                res.emplace_back(node->val);
                return;
            }

            depthFirstSearch(node->left, K - 1, parents, visited, res);
            depthFirstSearch(node->right, K - 1, parents, visited, res);
            depthFirstSearch(parents[node], K - 1, parents, visited, res);
        }
};

Solution 2

// The following block might slightly improve the execution time;
// Can be removed;
static const auto __optimize__ = []() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    return 0;
}();

// Most of headers are already included;
// Can be removed;
#include <cstdint>
#include <vector>
#include <unordered_map>
#include <unordered_set>

using ValueType = int;

static const struct Solution {
        const std::vector<ValueType> distanceK(
            TreeNode* root,
            TreeNode* target,
            ValueType K
        ) {
            getParent(root);
            depthFirstSearch(target, K);
            return res;
        }

    private:
        std::vector<ValueType> res;
        std::unordered_map<TreeNode*, TreeNode*> parents;
        std::unordered_set<TreeNode*> visited;

        const void getParent(
            TreeNode* node
        ) {
            if (!node) {
                return;
            }

            if (node->left) {
                parents[node->left] = node;
                getParent(node->left);
            }

            if (node->right) {
                parents[node->right] = node;
                getParent(node->right);
            }
        }



        const void depthFirstSearch(
            TreeNode* node,
            const ValueType K
        ) {
            if (!node) {
                return;
            }

            if (visited.count(node) > 0) {
                return;
            }

            visited.insert(node);

            if (!K) {
                res.emplace_back(node->val);
                return;
            }

            depthFirstSearch(node->left, K - 1);
            depthFirstSearch(node->right, K - 1);
            depthFirstSearch(parents[node], K - 1);
        }
};

Reference

Here is LeetCode's template:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        
    }
};

• 863. All Nodes Distance K in Binary Tree - Problem

• 863. All Nodes Distance K in Binary Tree - Discuss

Answer 1

About __optimize__

Identifiers starting with double underscores are reserved. Also, why is this written as a lambda instead of as a regular function? You're also not reading and writing to standard I/O, so this function wouldn't have any effect anyway.

Avoid creating type aliases outside of a class or namespace

Don't declare using ValueType = int in the global namespace, as it is a very generic name and could conflict with other code that does the same. In this case, just declare this inside struct Solution.

static const has no effect on a struct definition

The qualifiers static and const have no effect on a definition of a struct. It is allowed in the C++ grammar because you can define a struct and declare a variable of that type in one go, like:

static const struct foo {
    ...
} bar;

foo baz;

In the above, bar is static const, but baz is not.

const has no effect on non-pointer/reference return values

Likewise, const has no effect on the return value of a function, unless it's a pointer or reference that is returned. It doesn't make sense, because you are always allowed to copy a const value into a non-const variable. Also, what do you expect const void to mean?

Static vs. non-static member functions

The two variants differ in whether they use non-static member functions, with state kept as class member variables, or static member functions with state allocated on the stack and passed as pointers to other member functions. Both are valid approaches, although the fact that it doesn't really matter should tell you that the use of a struct Solution itself is pointless. In a real world application, you would have a function distanceK() that is not a member of any class. I believe LeetCode just gives you a class because they copy&pasted Java problems to C++ with minimal changes, and Java doesn't allow functions to be defined outside of a class.

A compiler, when optimization is enabled, will probably generate very similar assembly in both cases.

Use a std::bitset to keep track of visited nodes

The problem says that there are only up to 501 unique nodes. That means you can use a std::bitset to keep track of which nodes you visited. The bitset will only use 64 bytes, compare that to 8 bytes for a single TreeNode *, let alone all the other overhead of keeping an std::unordered_set<TreeNode *>.

Try to avoid using a lot of memory

An issue with your algorithm, which otherwise looks very reasonable, is that you need to calculate the parents of all the nodes. Since you cannot store it in a TreeNode itself, you now have to keep an unordered_map<TreeNode *, TreeNode *>, which takes up roughly as much space as the input.

If you do a depth-first search, then when calling the DFS function recursively, you know the parent of the node you're recursing, so you can pass that to the function, like so:

void DFS(Node *node, Node *parent) {
    if (!node)
        return;

    // do something with node and/or parent

    DFS(node->left, node);
    DFS(node->right, node);
}

The problem for you is that you want to start the DFS at the target node instead of at the root of the tree, so you don't know the parents of the target. However, you might be able to modify your algorithm to start from the root anyway, and then keep track of how far you had to descend to reach the target. Once you reach the target, you recurse down as usual, but when you're done you return upwards, where you should somehow signal that you've encountered the target, and then find nodes at distance K the other way. This might mean you have to visit parts of the tree twice though, but you already do that in your current algorithm.