[LeetCode] 863. All Nodes Distance K in Binary Tree

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/description/

All Nodes Distance K in Binary Tree - LeetCode

문제 유형 : BFS / DFS

문제 난이도 : Medium

 

문제

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

 

이진 트리의 루트가 주어진다. 타깃 노드 target이 주어지고, 정수 K가 주어진다.

target으로부터 k만큼의 거리에 있는 노드들의 값을 반환하여라.

순서는 무관하다.

 

풀이

두 번의 BFS를 진행한다 :

1. 이진 트리의 노드간의 부모를 정의하기 위한 BFS

2. target에서부터 k만큼의 거리에 있는 노드들을 찾기 위한 BFS

 

unordered_map을 이용하여, 노드의 값이 가지는 부모를 따로 저장한다.

 

코드(C++)

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        vector<int> nodes;
        unordered_map<int, TreeNode*> parent;
        queue<TreeNode*> q;
        q.push(root);
		// 전체 노드를 탐색하며 부모 관계를 잇는다.
        while(!q.empty()){
            int size = q.size();
            for(int i = 0; i < size; i++){
                auto curr = q.front();
                q.pop();
                //왼쪽에 값이 있는 경우, 현재 노드의 왼쪽의 부모를 현재 노드로 하고, 큐에 넣는다.
                if(curr -> left){
                    parent[curr->left->val] = curr;
                    q.push(curr->left); 
                }
				//오른쪽에 값이 있는 경우, 현재 노드의 오른쪽의 부모를 현재 노드로 하고, 큐에 넣는다.
                if(curr -> right){
                    parent[curr->right->val] = curr;
                    q.push(curr->right); 
                }
            }
        }
		//k만큼 떨어져있는 노드들 찾기. 이 while문이 끝나면 큐에는 k만큼 떨어져있는 노드들만 남음.
        unordered_map<int, int> visited;
        q.push(target);
        while(k-- && !q.empty()){
            int size = q.size();

            for(int i = 0; i < size; i++){
                auto curr = q.front();
                q.pop();

                visited[curr -> val] = 1;

                if(curr -> left && !visited[curr->left->val]){
                    q.push(curr -> left);
                }

                if(curr -> right && !visited[curr->right->val]){
                    q.push(curr -> right);
                }

                if(parent[curr->val] && !visited[parent[curr->val] -> val]){
                    q.push(parent[curr->val]);
                }

            }
        }
		//노드들 출력하기
        while(!q.empty()){
            nodes.push_back(q.front()->val);
            q.pop();
        }
        return nodes;
    }
};