from typing import (
    List,
)
from lintcode import (
    TreeNode,
)

"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None
"""

class Solution:
    """
    @param root: the given BST
    @param target: the given target
    @param k: the given k
    @return: k values in the BST that are closest to the target
             we will sort your return value in output
    """
    def closest_k_values(self, root: TreeNode, target: float, k: int) -> List[int]:
        if not root or k == 0:
            return []

        # write your code here
        lower_stack = self.get_stack(root, target) # init to found target
        upper_stack = list(lower_stack)

        # delete duplicate closest element, change wrong lower[-1] or upper[-1]
        if lower_stack[-1].val < target:
            self.move_upper(upper_stack)
        else:
            self.move_lower(lower_stack)

        ans = []
        for i in range(k):
            if self.is_lower_closer(lower_stack, upper_stack, target):
                ans.append(lower_stack[-1].val)
                self.move_lower(lower_stack)
            else:
                ans.append(upper_stack[-1].val)
                self.move_upper(upper_stack)
        return ans

    def is_lower_closer(self, lower_stack, upper_stack, target):
        if not lower_stack: # has nothing
            return False
        if not upper_stack:
            return True

        return target - lower_stack[-1].val < upper_stack[-1].val - target

    def move_lower(self, stack):
        if stack[-1].left: # has left, move to the right bottom
            node = stack[-1].left
            while node:
                stack.append(node)
                node = node.right
        else: # no left, move to the first turn right parent
            node = stack.pop()
            while stack and stack[-1].left == node:
                node = stack.pop()

    def move_upper(self, stack):
        if stack[-1].right: # has right, move to the left bottom
            node = stack[-1].right
            while node:
                stack.append(node)
                node = node.left
        else: # no right, move to the first turn left parent
            node = stack.pop()
            while stack and stack[-1].right == node:
                node = stack.pop()

    def get_stack(self, root, target):
        stack = []
        while root:
            stack.append(root)
            if target <= root.val:
                root = root.left
            else:
                root = root.right

        return stack