All Nodes Distance K in Binary Tree Solutions in C++

/// Source : https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/description/
/// Author : liuyubobobo
/// Time   : 2018-07-01

#include <iostream>
#include <vector>
#include <cassert>
#include <queue>

using namespace std;

/// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

/// Construct a Graph and Using BFS
/// Time Complexity: O(n)
/// Space Complexity: O(n)
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {

        if(K == 0)
            return {target->val};

        vector<vector<int>> g;
        vector<int> value;

        int targetIndex = -1;
        constructG(g, value, root, -1, target, targetIndex);
        assert(targetIndex >= 0);

        int n = value.size();
        assert(n == g.size());

        vector<int> res;
        vector<bool> visited(n, false);
        queue<pair<int, int>> q;
        q.push(make_pair(targetIndex, 0));
        visited[targetIndex] = true;
        while(!q.empty()){
            int curIndex = q.front().first;
            int curDis = q.front().second;
            q.pop();
            if(curDis == K)
                res.push_back(value[curIndex]);

            for(int nextIndex: g[curIndex])
                if(!visited[nextIndex]){
                    q.push(make_pair(nextIndex, curDis + 1));
                    visited[nextIndex] = true;
                }
        }
        return res;
    }

private:
    void constructG(vector<vector<int>>& g, vector<int>& value,
                    TreeNode* root, int parent, TreeNode* target, int& targetIndex){

        value.push_back(root->val);
        g.push_back(vector<int>());
        int index = value.size() - 1;

        if(parent != -1)
            g[index].push_back(parent);

        if(root->left != NULL){
            g[index].push_back(value.size());
            constructG(g, value, root->left, index, target, targetIndex);
        }

        if(root->right != NULL){
            g[index].push_back(value.size());
            constructG(g, value, root->right, index, target, targetIndex);
        }

        if(target == root)
            targetIndex = index;
    }
};


int main() {

    return 0;
}

/// Source : https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/description/
/// Author : liuyubobobo
/// Time   : 2018-07-01

#include <iostream>
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <cassert>
#include <queue>

using namespace std;

/// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

/// Construct a Graph and Using BFS
/// An mch easier way to construct the Graph :)
///
/// Time Complexity: O(n)
/// Space Complexity: O(n)
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {

        if(K == 0)
            return {target->val};

        unordered_map<TreeNode*, TreeNode*> parents;
        dfs(root, NULL, parents);

        vector<int> res;
        unordered_set<TreeNode*> visited;

        queue<pair<TreeNode*, int>> q;
        q.push(make_pair(target, 0));
        visited.insert(target);
        while(!q.empty()){
            TreeNode* curNode = q.front().first;
            int curDis = q.front().second;
            q.pop();
            if(curDis == K)
                res.push_back(curNode->val);

            unordered_map<TreeNode*, TreeNode*>::iterator piter = parents.find(curNode);
            if(piter != parents.end() && visited.find(piter->second) == visited.end()){
                q.push(make_pair(piter->second, curDis + 1));
                visited.insert(piter->second);
            }

            if(curNode->left != NULL && visited.find(curNode->left) == visited.end()){
                q.push(make_pair(curNode->left, curDis + 1));
                visited.insert(curNode->left);
            }

            if(curNode->right != NULL && visited.find(curNode->right) == visited.end()){
                q.push(make_pair(curNode->right, curDis + 1));
                visited.insert(curNode->right);
            }
        }
        return res;
    }

private:
    void dfs(TreeNode* root, TreeNode* parent,
             unordered_map<TreeNode*, TreeNode*>& parents){

        if(root == NULL)
            return;

        if(parent != NULL)
            parents[root] = parent;

        dfs(root->left, root, parents);
        dfs(root->right, root, parents);
    }
};


int main() {

    return 0;
}

/// Source : https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/description/
/// Author : liuyubobobo
/// Time   : 2018-07-01

#include <iostream>
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <cassert>
#include <queue>

using namespace std;

/// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

/// DFS and get the results during the recursion
///
/// Time Complexity: O(n)
/// Space Complexity: O(n)
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {

        if(K == 0)
            return {target->val};

        vector<int> res;
        assert(dfs(root, target, K, res) != -1);
        return res;
    }

private:
    int dfs(TreeNode* root, TreeNode* target, int K, vector<int>& res){

        if(root == NULL)
            return -1;

        if(root == target){
            addRes(root, K, res);
            return 0;
        }

        int L = dfs(root->left, target, K, res);
        if(L != -1){
            if(L + 1 == K)
                addRes(root, 0, res);
            else
                addRes(root->right, K - L - 2, res);
            return L + 1;
        }

        int R = dfs(root->right, target, K, res);
        if(R != -1){
            if(R + 1 == K)
                addRes(root, 0, res);
            else
                addRes(root->left, K - R - 2, res);
            return R + 1;
        }

        return -1;
    }

    void addRes(TreeNode* root, int dis, vector<int>& res){

        if(root == NULL)
            return;

        if(dis == 0){
            res.push_back(root->val);
            return;
        }

        addRes(root->left, dis - 1, res);
        addRes(root->right, dis - 1, res);
    }
};


int main() {

    return 0;
}