package com.colbert.binarytree

import com.sun.jmx.remote.internal.ArrayQueue
import java.util.*

class BinaryTree(var rootNode: TreeNode = BinaryTree.TreeNode(1, "A")) {

    class TreeNode {
        var index: Int
        var data: String
        var leftChild: TreeNode?
        var rightChild: TreeNode?

        constructor(index: Int, data: String) {
            this.index = index
            this.data = data
            this.leftChild = null
            this.rightChild = null
        }
    }

    /**
     * 构建二叉树
     *             A
     *       B          C
     *  D        E F         G
     */
    fun createBinaryTree() {
        var nodeB = TreeNode(2, "B")
        var nodeC = TreeNode(3, "C")
        var nodeD = TreeNode(4, "D")
        var nodeE = TreeNode(5, "E")
        var nodeF = TreeNode(6, "F")
        var nodeG = TreeNode(6, "G")
        rootNode.leftChild = nodeB;
        rootNode.rightChild = nodeC
        nodeB.leftChild = nodeD
        nodeB.rightChild = nodeE
        nodeC.leftChild = nodeF
        nodeC.rightChild = nodeG
    }

    /**
     * 树的高度
     */
    fun getHeight(): Int {
        return getHeight(rootNode);
    }

    fun getHeight(node: TreeNode?): Int {
        if (node == null) {
            return 0
        } else {
            return Math.max(getHeight(node?.leftChild), getHeight(node?.rightChild)) + 1
        }
    }

    /**
     * 树的节点数
     */
    fun getSize(): Int {
        return getSize(rootNode)
    }

    fun getSize(node: TreeNode?): Int {
        if (node == null) {
            return 0
        } else {
            // 父节点(1) + 左节点 + 右节点
            return getSize(node?.leftChild) + getSize(node?.rightChild) + 1
        }
    }

    /**
     * 前序遍历
     * 顺序：先遍历根节点，再遍历左子节点，最后遍历右子节点
     */
    fun preOrder() {
        preOrder(rootNode)
    }

    fun preOrder(node: TreeNode?) {
        if (node == null) return
        print(node?.data + " ")
        preOrder(node?.leftChild)
        preOrder(node?.rightChild)
    }

    /**
     * 前序遍历
     * 非迭代实现 （栈）
     */
    fun stackPreOrder() {
        if (rootNode == null) return
        var stack: Stack<TreeNode> = Stack()
        stack.push(rootNode)
        do {
            var node: TreeNode = stack.pop()
            print(node.data + " ")
            if (node.rightChild != null) {
                stack.push(node.rightChild)
            }
            if (node.leftChild != null) {
                stack.push(node.leftChild)
            }
        } while (!stack.isEmpty())
    }

    /**
     * 中序遍历
     * 顺序：先遍历左子节点，再遍历根节点，最后遍历右子节点
     */
    fun midOrder() {
        midOrder(rootNode)
    }

    fun midOrder(node: TreeNode?) {
        if (node == null) return
        midOrder(node?.leftChild)
        print(node?.data + " ")
        midOrder(node?.rightChild)
    }

    /**
     * 中序遍历
     * 非迭代实现 （栈）
     */
    fun stackMidOrder() {
        if (rootNode == null) return
        var stack: Stack<TreeNode> = Stack()
        var node: TreeNode? = rootNode
        while (node != null || !stack.isEmpty()) {
            while (node != null) {
                stack.push(node)
                node = node?.leftChild
            }
            node = stack.pop()
            print(node?.data + " ")
            node = node?.rightChild
        }
    }

    /**
     * 后序遍历
     * 顺序：先遍历左子节点，再遍历右子节点，最后遍历根节点
     */
    fun afterOrder() {
        afterOrder(rootNode)
    }

    fun afterOrder(node: TreeNode?) {
        if (node == null) return
        afterOrder(node?.leftChild)
        afterOrder(node?.rightChild)
        print(node?.data + " ")
    }

    /**
     * 后序遍历
     * 非迭代实现 （栈）
     */
    fun stackAfterOrder() {
        if (rootNode == null) return
        var node: TreeNode? = rootNode
        var stack: Stack<TreeNode?> = Stack()
        // 存放右节点是否被访问过
        var map = HashMap<TreeNode?, Boolean>();
        while (node != null || stack.isNotEmpty()) {
            while (node != null) {
                stack.push(node)
                node = node?.leftChild
            }
            node = stack.peek()
            if (map.getOrDefault(node, false)) {
                // 如果右节点被访问过，就打印出来，并把它弹出栈
                print(node?.data + " ")
                map.remove(stack.pop())
                node = null
            } else {
                map.put(node, true)
                node = node?.rightChild
            }
        }
    }

    /**
     * 层序遍历 （队列）
     * 顺序：按层次依次遍历树的节点
     */
    fun queueLevelOrder() {
        if (rootNode == null) return
        var queue = ArrayDeque<TreeNode>()
        queue.offer(rootNode)
        while (queue.isNotEmpty()) {
            // 将对头的元素弹出队列
            var node = queue.poll()
            print(node?.data + " ")
            if (node?.leftChild != null) {
                // 将左节点加入到对尾
                queue.offer(node?.leftChild)
            }
            if (node?.rightChild != null) {
                // 将右节点加入到对尾
                queue.offer(node?.rightChild)
            }
        }
    }
}

fun main(args: Array<String>) {
    var tree = BinaryTree()
    tree.createBinaryTree()

    // 树的高度
    print("树的高度 " + tree.getHeight())
    // 树的节点数
    print("\n树的节点数 " + tree.getSize())
    // 前序遍历
    print("\n前序遍历 ")
    tree.preOrder()
    // 前序栈遍历
    print("\n前序栈遍历 ")
    tree.stackPreOrder()
    // 中序遍历
    print("\n中序遍历 ")
    tree.midOrder()
    // 中序栈遍历
    print("\n中序栈遍历 ")
    tree.stackMidOrder()
    // 后序遍历
    print("\n后序遍历 ")
    tree.afterOrder()
    print("\n后序栈遍历 ")
    tree.stackAfterOrder()
    // 层序遍历
    print("\n层序遍历 ")
    tree.queueLevelOrder()
}