Binary Tree: Given a binary tree, connect nodes at the same level

Problem Statement:

Given a binary tree, you need to connect the nodes.

Each nodes will have a next pointer, set these pointers to the next right of each node.

Example:

Input: 

    /*
     *            4      ----> NULL
     *         /     \
     *        2  ->   6 -----> NULL
     *       / \      /  \
     *      1 ->3 -> 5 -> 7 -> NULL
     */

Solution Explanation:

We will use level order traversal.
When each nodes are processed, set each nodes nect_right to the next node in the queue.

Time Complexity: O(n)
Space Complexity: O(n)

Code Solution

#include <iostream>
#include <vector>
#include <queue>
using namespace std;

struct Node 
{
    int data;
    struct Node *left;
    struct Node *right;
    struct Node *next_right;
};

void solution (Node* root) 
{

    queue<Node*> q;
        q.push(root);

    while (!q.empty()) 
    {

        int nodeCount = q.size();

        for(int i = 0; i < nodeCount; i++) 
        {
            Node* node = q.front();
            q.pop();

            // Set next_right
            if (i == nodeCount - 1)
                node->next_right = nullptr;
            else
                node->next_right = q.front();

            // insert children
            if (node->left) q.push(node->left);
            if (node->right) q.push(node->right);
        }
    }
}

struct Node* newNode(int data)
{
    struct Node* newNode = new Node;
    newNode->data = data;
    newNode->left = newNode->right = NULL;

    return (newNode);
}

void display_inorder(Node* root)
{
    if (root == NULL)
        return;

    display_inorder(root->left);
    cout << root->data << " ";
    display_inorder(root->right);
}



void insert_inorder(Node* root, vector<int>& nodes) 
{
    if (root == nullptr) 
    {
        return;
    }
  
    insert_inorder(root->left, nodes);  
    nodes.push_back(root->data);          
    insert_inorder(root->right, nodes); 
}

int main(void)
{
    /*
     *           10
     *         /    \
     *        8      12
     *       / \    /  \
     *      2   9  11   14
     */
    struct Node* root = newNode(10);
    root->left = newNode(8);
    root->right = newNode(12);
    root->left->left = newNode(2);
    root->left->right = newNode(9);
    root->right->left = newNode(11);
    root->right->right = newNode(14);

    solution(root);

    cout << "Right of 8 is " << root->left->next_right->data << endl;
    cout << "Right of 9 is " << root->left->right->next_right->data << endl;

    return 0;
}

Output

Right of 8 is 12
Right of 9 is 11