剑指offer-二叉树层次遍历求树深度-python

# -*- coding:utf-8 -*-
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
class Solution:
    def TreeDepth(self, pRoot):
        # write code here
        # 层次遍历
        if not pRoot:
            return 0
        queue = [pRoot]
        count = 0
        total_p = 1
        depth = 0
        while queue:
            head = queue.pop(0)
            count+=1
            if head.left:
                queue.append(head.left)
            if head.right:
                queue.append(head.right)
            if count == total_p:
                count = 0
                total_p = len(queue)
                depth+=1
        return depth