Find the number of nodes in the widest layer of a binary tree

Which layer has the most nodes , It is the widest layer .

C++ edition

/************************************************************************* > File Name: 031. Find how many nodes there are in the widest layer of the binary tree .cpp > Author: Maureen > Mail: [email protected] > Created Time:  Two  6/21 19:45:56 2022 ************************************************************************/

#include <iostream>
#include <queue>
#include <ctime>
#include <unordered_map>
using namespace std;


class TreeNode {
    
public:
    int value;
    TreeNode *left;
    TreeNode *right;

    TreeNode(int v) : value(v) {
    }
};

// Use a limited number of variables 
int maxWidth(TreeNode *root) {
    
    //cout << "maxWidth" << endl;
    if (root == nullptr) return 0;

    queue<TreeNode *> que;
    que.push(root);
    
    // The rightmost node of the current layer 
    TreeNode *curEnd = root;
    // The rightmost node of the next layer 
    TreeNode *nextEnd = nullptr;

    int ans = 0;
    // The number of nodes in the current layer 
    int curLevelNodeCnt = 0;
    
    // Sequence traversal 
    while (!que.empty()) {
    
        TreeNode *cur = que.front();
        //cout << "cur = " << cur << endl;
        que.pop();

        if (cur->left != nullptr) {
    
            que.push(cur->left);
            // When counting the current layer , Prepare for the next floor 
            nextEnd = cur->left;
        }

        if (cur->right != nullptr) {
    
            que.push(cur->right);
            nextEnd = cur->right;
        }
        
        // The number of nodes in the current layer ++
        curLevelNodeCnt++;
        if (cur == curEnd) {
    // Traverse to the last node of the current layer 
            ans = max(ans, curLevelNodeCnt);
            // Start the next level of traversal 
            curLevelNodeCnt = 0;
            curEnd = nextEnd;
        }
    }
    //cout << "maxWidth end!" << endl;
    return ans;
}

// Using a container unordered_map
int maxWidthUseMap(TreeNode *root) {
    
    if (root == nullptr) return 0;

    queue<TreeNode *> que;
    que.push(root);
    
    //key Is the node ,value On which floor ,
    unordered_map<TreeNode *, int> levelMap;
    levelMap[root] = 1;

    int curLevel = 1;
    int curLevelNodeCnt = 0;
    int ans = 0;

    while (!que.empty()) {
    
        TreeNode *cur = que.front();
        que.pop();

        int curNodeLevel = levelMap[cur];

        if (cur->left != nullptr) {
    
            levelMap[cur->left] = curNodeLevel + 1;
            que.push(cur->left);
        }

        if (cur->right != nullptr) {
    
            levelMap[cur->right] = curNodeLevel + 1;
            que.push(cur->right);
        }

        if (curNodeLevel == curLevel) {
    
            curLevelNodeCnt++;
        } else {
    
            ans = max(ans, curLevelNodeCnt);
            curLevel++;
            curLevelNodeCnt = 1;
        }
    }

    ans = max(ans, curLevelNodeCnt);

    return ans;
}

//For test 

TreeNode *generate(int level, int maxLevel, int maxValue) {
    
    if (level > maxLevel || rand() % 100 < 0.5) return nullptr;

    TreeNode *root = new TreeNode(rand() % maxValue);
    root->left = generate(level + 1, maxLevel, maxValue);
    root->right = generate(level + 1, maxLevel, maxValue);
    return root;
}

TreeNode *generateRandomBST(int maxLevel, int maxValue) {
    
    return generate(1, maxLevel, maxValue);
}


TreeNode  *test() {
    
    TreeNode *root = new TreeNode(1);
    root->left = new TreeNode(2);
    root->right = new TreeNode(3);
    root->left->left = new TreeNode(4);
    root->left->right = new TreeNode(5);
    root->right->left = new TreeNode(6);
    root->right->right = new TreeNode(7);
    return root;
}

int main() {
    
    srand(time(0));
    int maxLevel = 10;
    int maxValue = 100;
    int testTime = 1000000;
    cout << "test begin" << endl;
    for (int i = 0; i < testTime + 1; i++) {
    
        TreeNode *root = generateRandomBST(maxLevel, maxValue);
        //TreeNode *root = test();
        if (maxWidth(root) != maxWidthUseMap(root)) {
    
            cout << maxWidth(root) << ", " << maxWidthUseMap(root) << endl;
            cout << "Oops!" << endl;
            break;
        }

        if (i && i % 1000 == 0) cout << i << " cases passed!" << endl;
    }
    cout << "finish!" << endl;
    return 0;
}

Java edition

import java.util.HashMap;
import java.util.LinkedList;
import java.util.Queue;

public class TreeMaxWidth {
    

	public static class Node {
    
		public int value;
		public Node left;
		public Node right;

		public Node(int data) {
    
			this.value = data;
		}
	}

	public static int maxWidthUseMap(Node head) {
    
		if (head == null) {
    
			return 0;
		}
		Queue<Node> queue = new LinkedList<>();
		queue.add(head);
		// key  stay   Which floor ,value
		HashMap<Node, Integer> levelMap = new HashMap<>();
		levelMap.put(head, 1);
		int curLevel = 1; //  At present, you are counting the width of which floor 
		int curLevelNodes = 0; //  Current layer curLevel layer , What is the current width 
		int max = 0;
		while (!queue.isEmpty()) {
    
			Node cur = queue.poll();
			int curNodeLevel = levelMap.get(cur);
			if (cur.left != null) {
    
				levelMap.put(cur.left, curNodeLevel + 1);
				queue.add(cur.left);
			}
			if (cur.right != null) {
    
				levelMap.put(cur.right, curNodeLevel + 1);
				queue.add(cur.right);
			}
			if (curNodeLevel == curLevel) {
    
				curLevelNodes++;
			} else {
    
				max = Math.max(max, curLevelNodes);
				curLevel++;
				curLevelNodes = 1;
			}
		}
		max = Math.max(max, curLevelNodes);
		return max;
	}

	public static int maxWidthNoMap(Node head) {
    
		if (head == null) {
    
			return 0;
		}
		Queue<Node> queue = new LinkedList<>();
		queue.add(head);
		Node curEnd = head; //  Current layer , Who is the rightmost node 
		Node nextEnd = null; //  The next layer , Who is the rightmost node 
		int max = 0;
		int curLevelNodes = 0; //  The number of nodes in the current layer 
		while (!queue.isEmpty()) {
    
			Node cur = queue.poll();
			if (cur.left != null) {
    
				queue.add(cur.left);
				nextEnd = cur.left;
			}
			if (cur.right != null) {
    
				queue.add(cur.right);
				nextEnd = cur.right;
			}
			curLevelNodes++;
			if (cur == curEnd) {
    
				max = Math.max(max, curLevelNodes);
				curLevelNodes = 0;
				curEnd = nextEnd;
			}
		}
		return max;
	}

	// for test
	public static Node generateRandomBST(int maxLevel, int maxValue) {
    
		return generate(1, maxLevel, maxValue);
	}

	// for test
	public static Node generate(int level, int maxLevel, int maxValue) {
    
		if (level > maxLevel || Math.random() < 0.5) {
    
			return null;
		}
		Node head = new Node((int) (Math.random() * maxValue));
		head.left = generate(level + 1, maxLevel, maxValue);
		head.right = generate(level + 1, maxLevel, maxValue);
		return head;
	}

	public static void main(String[] args) {
    
		int maxLevel = 10;
		int maxValue = 100;
		int testTimes = 1000000;
		for (int i = 0; i < testTimes; i++) {
    
			Node head = generateRandomBST(maxLevel, maxValue);
			if (maxWidthUseMap(head) != maxWidthNoMap(head)) {
    
				System.out.println("Oops!");
			}
		}
		System.out.println("finish!");

	}
}