raghuvanshraj

nodes-at-given-distance-in-binary-tree.cpp

Jul 29th, 2021
795
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. // { Driver Code Starts
  2. #include<bits/stdc++.h>
  3. using namespace std;
  4.  
  5. // Tree Node
  6. struct Node
  7. {
  8.     int data;
  9.     Node* left;
  10.     Node* right;
  11. };
  12.  
  13. // Utility function to create a new Tree Node
  14. Node* newNode(int val)
  15. {
  16.     Node* temp = new Node;
  17.     temp->data = val;
  18.     temp->left = NULL;
  19.     temp->right = NULL;
  20.    
  21.     return temp;
  22. }
  23.  
  24. // Function to Build Tree
  25. Node* buildTree(string str)
  26. {  
  27.     // Corner Case
  28.     if(str.length() == 0 || str[0] == 'N')
  29.             return NULL;
  30.    
  31.     // Creating vector of strings from input
  32.     // string after spliting by space
  33.     vector<string> ip;
  34.    
  35.     istringstream iss(str);
  36.     for(string str; iss >> str; )
  37.         ip.push_back(str);
  38.        
  39.     // Create the root of the tree
  40.     Node* root = newNode(stoi(ip[0]));
  41.        
  42.     // Push the root to the queue
  43.     queue<Node*> queue;
  44.     queue.push(root);
  45.        
  46.     // Starting from the second element
  47.     int i = 1;
  48.     while(!queue.empty() && i < ip.size()) {
  49.            
  50.         // Get and remove the front of the queue
  51.         Node* currNode = queue.front();
  52.         queue.pop();
  53.            
  54.         // Get the current node's value from the string
  55.         string currVal = ip[i];
  56.            
  57.         // If the left child is not null
  58.         if(currVal != "N") {
  59.                
  60.             // Create the left child for the current node
  61.             currNode->left = newNode(stoi(currVal));
  62.                
  63.             // Push it to the queue
  64.             queue.push(currNode->left);
  65.         }
  66.            
  67.         // For the right child
  68.         i++;
  69.         if(i >= ip.size())
  70.             break;
  71.         currVal = ip[i];
  72.            
  73.         // If the right child is not null
  74.         if(currVal != "N") {
  75.                
  76.             // Create the right child for the current node
  77.             currNode->right = newNode(stoi(currVal));
  78.                
  79.             // Push it to the queue
  80.             queue.push(currNode->right);
  81.         }
  82.         i++;
  83.     }
  84.    
  85.     return root;
  86. }
  87.  
  88.  
  89.  // } Driver Code Ends
  90. /* A binary Tree node
  91. struct Node
  92. {
  93.     int data;
  94.     struct Node *left, *right;
  95. };
  96. */
  97.  
  98. class Solution
  99. {
  100. private:
  101.     vector<int> ans;
  102. public:
  103.     void _KDistanceNodesDown(Node *root, int k) {
  104.         if (!root || k < 0) {
  105.             return;
  106.         }
  107.        
  108.         if (k == 0) {
  109.             ans.push_back(root->data);
  110.         }
  111.        
  112.         _KDistanceNodesDown(root->left, k-1);
  113.         _KDistanceNodesDown(root->right, k-1);
  114.     }
  115.    
  116.     int _KDistanceNodes(Node *root, int target, int k) {
  117.         if (!root) {
  118.             return 0;
  119.         }
  120.        
  121.         if (root->data == target) {
  122.             _KDistanceNodesDown(root, k);
  123.             return 1;
  124.         }
  125.        
  126.         int left = _KDistanceNodes(root->left, target, k);
  127.         if (left != 0) {
  128.             if (left == k) {
  129.                 ans.push_back(root->data);
  130.             } else {
  131.                 _KDistanceNodesDown(root->right, k - left - 1);
  132.             }
  133.             return left + 1;
  134.         }
  135.        
  136.         int right = _KDistanceNodes(root->right, target, k);
  137.         if (right != 0) {
  138.             if (right == k) {
  139.                 ans.push_back(root->data);
  140.             } else {
  141.                  _KDistanceNodesDown(root->left, k - right - 1);
  142.             }
  143.             return right + 1;
  144.         }
  145.        
  146.         return 0;
  147.     }
  148.    
  149.     vector<int> KDistanceNodes(Node* root, int target , int k) {
  150.         ans.clear();
  151.         _KDistanceNodes(root, target, k);
  152.         sort(ans.begin(), ans.end());
  153.         return ans;
  154.     }
  155. };
  156.  
  157. // { Driver Code Starts.
  158.  
  159. int main()
  160. {
  161.     int t;
  162.     cin>>t;
  163.     getchar();
  164.    
  165.     Solution x = Solution();
  166.    
  167.     while(t--)
  168.     {
  169.         string s;
  170.         getline(cin,s);
  171.         Node* head = buildTree(s);
  172.        
  173.         int target, k;
  174.         cin>> target >> k;
  175.         getchar();
  176.        
  177.         vector <int> res = x.KDistanceNodes(head, target, k);
  178.        
  179.         for( int i=0; i<res.size(); i++ )
  180.             cout<< res[i] << " ";
  181.         cout<<endl;
  182.     }
  183.     return 0;
  184. }
  185.   // } Driver Code Ends
RAW Paste Data