inorder_traverse_coroutine.py

class Node(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

    def __str__(self):
        return "(%s|%s|%s)" % (
            '*' if self.left else '_',
            self.val,
            '*' if self.right else '_')


def inorder_traverse(root):
    """
    In theory, any recursive function can be converted
    into a non-recursive function by using a manually
    managed stack to simulate the call stack

    think about what the recursive version looks like
    >>> def inorder(tree):
        ....if tree.left:
        ....    inorder(tree.left)
        ....
        ....do(tree.val)
        ....
        ....if tree.right:
        ....    inorder(tree.right)
    """
    stack = []
    # stack contains (node, processed)
    # where the processed means whether the node has been travesed

    def push_to_end(node):
        while node:
            stack.append((node, False))
            node = node.left

    push_to_end(root)

    while stack:
        node, processed = stack.pop()
        if processed:
            continue

        # do on node
        yield node

        if node.right:
            stack.append((node, True))
            push_to_end(node.right)


def inorder_inc_from(root, x):
    stack = []
    def push_to_end(node):
        while node:
            stack.append(node)
            node = node.left

    node = root
    while node:
        if node.val < x:
            # this is redundant
            #stack.append((node, True))
            node = node.right
            continue

        # node.val <= x
        stack.append(node)
        if not node.left:
            break
        else:
            node = node.left

    while stack:
        node = stack.pop()

        # do on node
        yield node.val

        if node.right:
            push_to_end(node.right)


def inorder_dec_from(root, x):
    stack = []
    def push_to_end(node):
        while node:
            stack.append(node)
            node = node.right

    node = root
    while node:
        if node.val >= x:
            node = node.left
            continue

        # node.val < x
        stack.append(node)
        if not node.right:
            break
        else:
            node = node.right

    while stack:
        node = stack.pop()

        # do on node
        yield node.val

        if node.left:
            push_to_end(node.left)


def yield_nearest(root, x):
    inc_gen = inorder_inc_from(root, x)
    dec_gen = inorder_dec_from(root, x)
    inc_next = next(inc_gen, None)
    dec_next = next(dec_gen, None)
    while inc_next != None or dec_next != None:
        if inc_next == None:
            yield dec_next
            for y in dec_gen:
                yield y
            break
        elif dec_next == None:
            yield inc_next
            for y in inc_gen:
                yield y
            break
        else:
            if abs(inc_next - x) < abs(dec_next - x):
                yield inc_next
                inc_next = next(inc_gen, None)
            else:
                yield dec_next
                dec_next = next(dec_gen, None)


def find_nearest_k(root, x, k):
    gen = yield_nearest(root, x)
    result = []
    while k > 0:
        result.append(next(gen))
        k -= 1
    return result