Zig-Zag level order traversal of Binary Tree after every K levels

Given a binary tree and an integer k, the task is to print the level order traversal in such a way that first k levels are printed from left to right, next k levels are printed from right to left, then next k levels are from left to right and so on.

Examples:

Input: k = 1

Output:
1
3 2
4 5 6
Explanation: In the above example, first level is printed from left to right and the second level is printed from right to left, and then last level is printed from left to right.
Input: k = 3

Output:
20
8 22
4 12
14 10
Explanation: In the above example, first 3 levels are printed from left to right and the last level is printed from right to left

Approach:

The idea is to start performing level order traversal on the tree from the left most end. After every k level, change the direction of printing the elements.
For this use stack. When the levels are to be printed from the right keep those values in stack and print the stack elements one by one from top. Because of the last in first out property of stack the elements would printed in reverse order.

Follow the steps mentioned below to implement the above idea:

Use queue to perform level order traversal.
In each level:
- If it is to be printed from the left, print them and push their child nodes in the queue.
- If this level is to be printed from the right side push the elements in a stack and print them after the whole level is traversed.
- If k levels are covered change the direction of printing from the next one.

C++

// C++ implementation to print ZigZag level order
// after every k Levels
#include <bits/stdc++.h>
using namespace std;

class Node {
public:
    int data;
    Node* left;
    Node* right;

    Node(int x) {
        data = x;
        left = nullptr;
        right = nullptr;
    }
};

// Function to traverse the tree and store nodes
// in a 2D vector by reversing the direction of traversal
// after every n levels
vector<vector<int>> traverse(Node* root, int n) {
    vector<vector<int>> result;

    if (!root) return result;

    // Queue for level order traversal
    queue<Node*> q;

    // Stack for reverse level order traversal
    stack<Node*> s;

    bool right2left = false;
    int count = 0;

    q.push(root);
    while (!q.empty()) {
        int size = q.size();
        count++;
      
        // To store nodes at the current level
        vector<int> levelNodes; 

        while (size--) {
            root = q.front();
            q.pop();

            // Store nodes from left to right
            if (right2left == false) {
                levelNodes.push_back(root->data);
            }
          
            // Push nodes into the stack for
          	// reverse order
            else {
                s.push(root);
            }

            if (root->left)
                q.push(root->left);
            if (root->right)
                q.push(root->right);
        }

        // If reversing, store nodes in
     	// reverse order
        if (right2left == true) {
            while (!s.empty()) {
                levelNodes.push_back(s.top()->data);
                s.pop();
            }
        }
      
        // Add the current level to the result
        result.push_back(levelNodes);  

        // Reverse the direction after n levels
        if (count == n) {
            right2left = !right2left;
            count = 0;  
        }
    }

    return result;
}

void print2DArray(vector<vector<int>>& arr) {
    for (auto row : arr) {
        for (int val : row) {
            cout << val << " ";
        }
        cout << endl;
    }
}

int main() {

    // Binary tree structure:
    //
    //        1
    //       / \
    //      2   3
    //     / \   \
    //    4   5   6

    Node *root = new Node(1);
    root->left = new Node(2);
    root->right = new Node(3);
    root->left->left = new Node(4);
    root->left->right = new Node(5);
    root->right->right = new Node(6);
    int k = 3;

    vector<vector<int>> result = traverse(root, k);
    print2DArray(result);

    return 0;
}

Java

// Java implementation to print ZigZag level order
// after every k Levels
import java.util.*;

class Node {
    int data;
    Node left, right;

    Node(int x) {
        data = x;
        left = null;
        right = null;
    }
}

class GfG {

    // Function to traverse the tree and store nodes
    // in a 2D List by reversing the direction of traversal
    // after every n levels
   	static List<List<Integer>> 
                       traverse(Node root, int n) {
      
        List<List<Integer>> result = new ArrayList<>();

        if (root == null) return result;

        // Queue for level order traversal
        Queue<Node> q = new LinkedList<>();

        // Stack for reverse level order traversal
        Stack<Node> s = new Stack<>();

        // For changing the direction of traversal
        boolean right2left = false;

        int count = 0;

        q.add(root);
        while (!q.isEmpty()) {
            int size = q.size();
            count++;
            
            // To store nodes at the current level
            List<Integer> levelNodes = new ArrayList<>();

            while (size-- > 0) {
                root = q.poll();

                // Store nodes from left to right
                if (!right2left) {
                    levelNodes.add(root.data);
                }
              
                // Push nodes into the stack for
              	// reverse order
                else {
                    s.push(root);
                }

                if (root.left != null) {
                    q.add(root.left);
                }
                if (root.right != null) {
                    q.add(root.right);
                }
            }

            // If reversing, store nodes
          	// in reverse order
            if (right2left) {
                while (!s.isEmpty()) {
                    levelNodes.add(s.pop().data);
                }
            }

            // Add the current level to the result
            result.add(levelNodes);

            // Reverse the direction after n levels
            if (count == n) {
                right2left = !right2left;
                count = 0;
            }
        }

        return result;
    }

    static void print2DArray(List<List<Integer>> arr) {
        for (List<Integer> row : arr) {
            for (int val : row) {
                System.out.print(val + " ");
            }
            System.out.println();
        }
    }

    public static void main(String[] args) {

        // Binary tree structure:
        //
        //        1
        //       / \
        //      2   3
        //     / \   \
        //    4   5   6
      
        Node root = new Node(1);
        root.left = new Node(2);
        root.right = new Node(3);
        root.left.left = new Node(4);
        root.left.right = new Node(5);
        root.right.right = new Node(6);
        int k = 3;
      
        List<List<Integer>> result = traverse(root, k);
        print2DArray(result);
    }
}

Python

# Python implementation to print ZigZag level order
# after every K Levels
from collections import deque

class Node:
    def __init__(self, x):
        self.data = x
        self.left = None
        self.right = None

# Function to traverse the tree and store nodes
# in a 2D list by reversing the direction of traversal
# after every n levels
def traverse(root, n):
    result = []

    if not root:
        return result

    # Queue for level order traversal
    q = deque()

    # Stack for reverse level 
    # order traversal
    s = []

    # For changing the direction of
    # traversal
    right2left = False

    count = 0

    q.append(root)
    while q:
        size = len(q)
        count += 1
        
        # To store nodes at the current level
        level_nodes = []

        while size > 0:
            root = q.popleft()
            size -= 1

            # Store nodes from left to right
            if not right2left:
                level_nodes.append(root.data)
                
            # Push nodes into the stack for
            # reverse order
            else:
                s.append(root)

            if root.left:
                q.append(root.left)
            if root.right:
                q.append(root.right)

        # If reversing, store nodes in 
        # reverse order
        if right2left:
            while s:
                level_nodes.append(s.pop().data)

        # Add the current level to the result
        result.append(level_nodes)

        # Reverse the direction after n levels
        if count == n:
            right2left = not right2left
            
            count = 0

    return result

def print2DArray(arr):
    for row in arr:
        for val in row:
            print(val, end=" ")
        print()

if __name__ == "__main__":
  
    # Binary tree structure:
    #
    #        1
    #       / \
    #      2   3
    #     / \   \
    #    4   5   6
    root = Node(1)
    root.left = Node(2)
    root.right = Node(3)
    root.left.left = Node(4)
    root.left.right = Node(5)
    root.right.right = Node(6)

    k = 3
    result = traverse(root, k)
    print2DArray(result)

// C# implementation to print ZigZag level order
// after every k Levels
using System;
using System.Collections.Generic;

class Node {
    public int data;
    public Node left;
    public Node right;

    public Node(int x) {
        data = x;
        left = null;
        right = null;
    }
}

class GfG {
  
    // Function to traverse the tree and store nodes
    // in a 2D list by reversing the direction of traversal
    // after every n levels
    static List<List<int>> Traverse(Node root, int n) {
        List<List<int>> result = new List<List<int>>();

        if (root == null) return result;

        // Queue for level order traversal
        Queue<Node> q = new Queue<Node>();

        // Stack for reverse level order traversal
        Stack<Node> s = new Stack<Node>();

        // For changing the direction of traversal
        bool right2left = false;

        int count = 0;

        q.Enqueue(root);
        while (q.Count > 0) {
            int size = q.Count;
            count++;

            // To store nodes at the current level
            List<int> levelNodes = new List<int>();

            while (size-- > 0) {
                root = q.Dequeue();

                // Store nodes from left to right
                if (!right2left) {
                    levelNodes.Add(root.data);
                } 
                else {
                    s.Push(root);
                }

                if (root.left != null) q.Enqueue(root.left);
                if (root.right != null) q.Enqueue(root.right);
            }

            // If reversing, store nodes in
          	// reverse order
            if (right2left) {
                while (s.Count > 0) {
                    levelNodes.Add(s.Pop().data);
                }
            }

            // Add the current level to the result
            result.Add(levelNodes);

            // Reverse the direction after n levels
            if (count == n) {
                right2left = !right2left;

                count = 0;
            }
        }

        return result;
    }

    static void Print2DArray(List<List<int>> arr) {
        foreach (var row in arr) {
            foreach (int val in row) {
                Console.Write(val + " ");
            }
            Console.WriteLine();
        }
    }


    static void Main() {
      
     	 // Binary tree structure:
        //
        //        1
        //       / \
        //      2   3
        //     / \   \
        //    4   5   6
      
        Node root = new Node(1);
        root.left = new Node(2);
        root.right = new Node(3);
        root.left.left = new Node(4);
        root.left.right = new Node(5);
        root.right.right = new Node(6);
        int k = 3;

        List<List<int>> result = Traverse(root, k);
        Print2DArray(result);
    }
}

JavaScript

// Javascript implementation to print ZigZag level order
// after every k Levels
class Node {
    constructor(x) {
        this.data = x;
        this.left = null;
        this.right = null;
    }
}

// Function to traverse the tree and store nodes
// in a 2D array by reversing the direction of traversal
// after every n levels
function traverse(root, n) {
    const result = [];

    if (!root) return result;

    // Queue for level order traversal
    const q = [];

    // Stack for reverse level order traversal
    const s = [];

    // For changing the direction of traversal
    let right2left = false;

    let count = 0;

    q.push(root);
    while (q.length > 0) {
        const size = q.length;
        count++;

        // To store nodes at the current level
        const levelNodes = [];

        for (let i = 0; i < size; i++) {
            const node = q.shift();

            // Store nodes from left to right
            if (!right2left) {
                levelNodes.push(node.data);
            } 
            else {
                s.push(node);
            }

            if (node.left) q.push(node.left);
            if (node.right) q.push(node.right);
        }

        // If reversing, store nodes in 
        // reverse order
        if (right2left) {
            while (s.length > 0) {
                levelNodes.push(s.pop().data);
            }
        }

        // Add the current level to the result
        result.push(levelNodes);

        // Reverse the direction after n levels
        if (count === n) {
            right2left = !right2left;
            
            count = 0;
        }
    }

    return result;
}

function print2DArray(arr) {
    for (const row of arr) {
        console.log(row.join(' '));
    }
}

// Binary tree structure:
// 
//        1
//       / \
//      2   3
//     / \   \
//    4   5   6

let root = new Node(1);
root.left = new Node(2);
root.right = new Node(3);
root.left.left = new Node(4);
root.left.right = new Node(5);
root.right.right = new Node(6)

const k = 3;
const result = traverse(root, k);
print2DArray(result);

Output

1 
2 3 
4 5 6

Time Complexity: O(n), where n is the total number of nodes in the tree. Each node is visited once during traversal.
Auxiliary Space : O(n), due to the space used by the queue and stack.

Zig-Zag level order traversal of Binary Tree after every K levels

Explore