#!/usr/bin/perl

# traverse( $tree, $func )
#
# Traverse $tree in order, calling $func() for each element.
#    in turn




sub traverse {
    my $tree = shift or return;   # skip undef pointers
    my $func = shift;

    traverse( $tree->{left}, $func );
    &$func( $tree );
    traverse( $tree->{right}, $func );
}

# $node = bal_tree_find( $tree, $val[, $cmp ] )
#
# Search $tree looking for a node that has the value $val.
# If provided, $cmp compares values instead of <=>.
#
# the return value:
#     $node points to the node that has value $val
#        or undef if no node has that value

sub bal_tree_find {
    my ($tree, $val, $cmp) = @_;
    my $result;

    while ( $tree ) {
        my $relation = defined $cmp
            ? $cmp->( $tree->{val}, $val )
		: $tree->{val} <=> $val;

        # Stop when the desired node is found.
        return $tree if $relation == 0;

        # Go down to the correct subtree.
        $tree = $relation < 0 ? $tree->{left} : $tree->{right};
    }

    # The desired node doesn't exist.
    return undef;
}

# ($tree, $node) = bal_tree_add( $tree, $val, $cmp )
#
# Search $tree looking for a node that has the value $val;
#    add it if it does not already exist.
# If provided, $cmp compares values instead of <=>.
#
# the return values:
#     $tree points to the (possibly new or changed) subtree that
#        has resulted from the add operation
#     $node points to the (possibly new) node that contains $val

sub bal_tree_add {
    my ($tree, $val, $cmp) = @_;
    my $result;

    # Return a new leaf if we fell off the bottom.
    unless ( $tree ) {
        $result = {
	    left   => undef,
	    right  => undef,
	    val    => $val,
	    height => 1
            };
        return( $result, $result );
    }

    my $relation = defined $cmp
        ? $cmp->( $tree->{val}, $val )
	    : $tree->{val} <=> $val;

    # Stop when the desired node is found.
    return ( $tree, $tree ) if $relation == 0;

    # Add to the correct subtree.
    if ( $relation < 0 ) {
        ($tree->{left}, $result)  =
            bal_tree_add ( $tree->{left}, $val, $cmp );
    } else {
        ($tree->{right}, $result) =
            bal_tree_add ( $tree->{right}, $val, $cmp );
    }

    # Make sure that this level is balanced, return the
    #    (possibly changed) top and the (possibly new) selected node.
    return ( balance_tree( $tree ), $result );
}

# ($tree, $node) = bal_tree_del( $tree, $val, $cmp )
#
# Search $tree looking for a node that has the value $val,
#    and delete it if it exists.
# If provided, $cmp compares values instead of <=>.
#
# the return values:
#     $tree points to the (possibly empty or changed) subtree that
#        has resulted from the delete operation
#     if found, $node points to the node that contains $val
#     if not found, $node is undef

sub bal_tree_del {
    # An empty (sub)tree does not contain the target.
    my $tree = shift or return (undef,undef);

    my ($val, $cmp) = @_;
    my $node;

    my $relation = defined $cmp
	? $cmp->($val, $tree->{val})
            : $val <=> $tree->{val};

    if ( $relation != 0 ) {
	# Not this node, go down the tree.
	if ( $relation < 0 ) {
	    ($tree->{left},$node) =
		bal_tree_del( $tree->{left}, $val, $cmp );
	} else {
	    ($tree->{right},$node) =
		bal_tree_del( $tree->{right}, $val, $cmp );
	}

	# No balancing required if it wasn't found.
	return ($tree,undef) unless $node;
    } else {
	# Must delete this node.  Remember it to return it,
	$node = $tree;

	# but splice the rest of the tree back together first
	$tree = bal_tree_join( $tree->{left}, $tree->{right} );

	# and make the deleted node forget its children (precaution
	# in case the caller tries to use the node).
	$node->{left} = $node->{right} = undef;
    }

    # Make sure that this level is balanced, return the
    #    (possibly changed) top and (possibly undef) selected node.
    return ( balance_tree($tree), $node );
}


# $tree = bal_tree_join( $left, $right );
#
# Join two trees together into a single tree.

sub bal_tree_join {
    my ($l, $r) = @_;

    # Simple case - one or both is null.
    return $l unless defined $r;
    return $r unless defined $l;

    # Nope - we've got two real trees to merge.
    my $top;

    if ( $l->{height} > $r->{height} ) {
	$top = $l;
	$top->{right} = bal_tree_join( $top->{right}, $r );
    } else {
	$top = $r;
	$top->{left} = bal_tree_join( $l, $top->{left} );
    }
    return balance_tree( $top );
}

# $tree = balance_tree( $tree )

sub balance_tree {
    # An empty tree is balanced already.
    my $tree = shift or return undef;

    # An empty link is height 0.
    my $lh = defined $tree->{left} && $tree->{left}{height};
    my $rh = defined $tree->{right} && $tree->{right}{height};

    # Rebalance if needed, return the (possibly changed) root.
    if ( $lh > 1+$rh ) {
        return swing_right( $tree );
    } elsif ( $lh+1 < $rh ) {
        return swing_left( $tree );
    } else {
        # Tree is either perfectly balanced or off by one.
        # Just fix its height.
        set_height( $tree );
        return $tree;
    }
}

# set_height( $tree )
sub set_height {
    my $tree = shift;
    my $p;
    # get heights, an undef node is height 0
    my $lh = defined ( $p = $tree->{left}  ) && $p->{height};
    my $rh = defined ( $p = $tree->{right} ) && $p->{height};
    $tree->{height} = $lh < $rh ? $rh+1 : $lh+1;
}

# $tree = swing_left( $tree )
#
# change       t        to          r         or         rl
#             / \                  / \                 /    \
#            l   r                t   rr              t      r
#               / \              / \                 / \    / \
#              rl  rr           l   rl              l  rll rlr rr
#             / \                  / \
#           rll rlr              rll rlr
#
# t and r must both exist.
# The second form is used if height of rl is greater than height of rr
# (since the first form would then lead to the height of t at least 2
# more than the height of rr).
#
# Changing to the second form is done in two steps, with first a
# move_right(r) and then a move_left(t), so it goes:
#
# change       t        to          t     and then to    rl
#             / \                  / \                 /    \
#            l   r                l   rl              t      r
#               / \                  / \             / \    / \
#              rl  rr              rll  r           l  rll rlr rr
#             / \                      / \
#           rll rlr                  rlr  rr

sub swing_left {
    my $tree = shift;
    my $r = $tree->{right};         # must exist
    my $rl = $r->{left};            # might exist
    my $rr = $r->{right};           # might exist
    my $l = $tree->{left};          # might exist

    # get heights, an undef node has height 0
    my $lh = $l && $l->{height};
    my $rlh = $rl && $rl->{height};
    my $rrh = $rr && $rr->{height};

    if ( $rlh > $rrh ) {
        $tree->{right} = move_right( $r );
    }

    return move_left( $tree );
}

# and the opposite swing

sub swing_right {
    my $tree = shift;
    my $l = $tree->{left};          # must exist
    my $lr = $l->{right};           # might exist
    my $ll = $l->{left};            # might exist
    my $r = $tree->{right};         # might exist

    # get heights, an undef node has height 0
    my $rh = $r && $r->{height};
    my $lrh = $lr && $lr->{height};
    my $llh = $ll && $ll->{height};

    if ( $lrh > $llh ) {
        $tree->{left} = move_left( $l );
    }

    return move_right( $tree );
}

# $tree = move_left( $tree )
#
# change       t        to          r
#             / \                  / \
#            l   r                t   rr
#               / \              / \
#              rl  rr           l   rl
#
# caller has determined that t and r both exist
#   (l can be undef, so can one of rl and rr)

sub move_left {
    my $tree = shift;
    my $r = $tree->{right};
    my $rl = $r->{left};

    $tree->{right} = $rl;
    $r->{left} = $tree;
    set_height( $tree );
    set_height( $r );
    return $r;
}

# $tree = move_right( $tree )
#
# opposite change from move_left

sub move_right {
    my $tree = shift;
    my $l = $tree->{left};
    my $lr = $l->{right};

    $tree->{left} = $lr;
    $l->{right} = $tree;
    set_height( $tree );
    set_height( $l );
    return $l;
}

# The tree starts out empty.
my $tree = undef;
my $node;

foreach ( 1..8 ) {
    ($tree, $node) = bal_tree_add( $tree, $_ * $_ );
}

($tree, $node) = bal_tree_del( $tree, 7*7 );