10
\$\begingroup\$

I'm trying to implement a left leaning red black tree as described here. This is their snippet for insert

private Node insert(Node h, Key key, Value value) {
    if (h == null) return new Node(key, value);

    if (isRed(h.left) && isRed(h.right)) colorFlip(h);

    int cmp = key.compareTo(h.key);
    if (cmp == 0) h.val = value;
    else if (cmp < 0) h.left = insert(h.left, key, value);
    else h.right = insert(h.right, key, value);

    if (isRed(h.right) && !isRed(h.left)) h = rotateLeft(h);
    if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);

    return h;
}

My attempt at a port in Haskell is... quite ugly, with a lot of repetition. I think it's because I'm still thinking procedurally. Any feedback on what I should do differently? Or is there no way around having many next-state variables (x', x'', x''')? Should I be approaching this completely differently?

data Colour = Red | Black deriving (Show)

data Tree a
    = Branch (Tree a) a (Tree a) Colour
    | Leaf
    deriving (Show)

add :: (Ord a) => Tree a -> a -> Tree a
add tree val
    = let
        (Branch left' node' right' _) = fix_up $ do_add tree val
    in (Branch left' node' right' Black) -- root always black

do_add :: (Ord a) => Tree a -> a -> Tree a
do_add (Branch left node right colour) val
    | val < node = (Branch (add left val) node right colour)
    | val > node = (Branch left node (add right val) colour)
    | otherwise = (Branch left node right colour)
do_add Leaf val = (Branch Leaf val Leaf Black)

get_left_node :: Tree a -> Tree a
get_left_node (Branch left _ _ _) = left
get_left_node Leaf = Leaf

fix_up :: Tree a -> Tree a
fix_up (Branch left node right colour)
    = let
        branch' = if ((not (is_red left)) && (is_red right)) then (rotate_left (Branch left node right colour)) else (Branch left node right colour)
        (Branch left' _ right' _) = branch'
        branch'' = if ((is_red left') && (is_red (get_left_node left'))) then (rotate_right branch') else branch'
        (Branch left'' _ right'' _) = branch''
        branch''' = if ((is_red left'') && (is_red right'')) then (flip_colours branch'') else branch''
    in branch'''


rotate_left :: Tree a -> Tree a
rotate_left (Branch left node (Branch right_left right_node right_right right_colour) colour)
    = let
        left' = (Branch left node right_left Red)
        centre' = (Branch left' right_node right_right colour)
    in centre'

rotate_right :: Tree a -> Tree a
rotate_right (Branch (Branch left_left left_node left_right left_colour) node right colour)
    = let
        right' = (Branch left_right node right Red)
        centre' = (Branch left_left left_node right' colour)
    in centre'

flip_colours :: Tree a -> Tree a
flip_colours (Branch (Branch left_left left_node left_right left_colour) node (Branch right_left right_node right_right right_colour) colour) = let
        left' = (Branch left_left left_node left_right (invert_colour left_colour))
        right' = (Branch right_left right_node right_right (invert_colour right_colour))
        centre' = (Branch left' node right' (invert_colour colour))
    in centre'

is_red :: Tree a -> Bool
is_red (Branch _ _ _ Red) = True
is_red _ = False

is_black :: Tree a -> Bool
is_black node = not $ is_red node
\$\endgroup\$
  • 1
    \$\begingroup\$ Could you please make the code snippet self-contained, add the data type definition etc., so that it can be readily compiled? It'd help analyzing and perhaps refactoring it a bit. \$\endgroup\$ – Petr Pudlák Jan 22 '14 at 19:53
  • \$\begingroup\$ @PetrPudlák sorry for the late reply. I added all the code. (I haven't tested it extensively but it seems to work for basic examples at in ghci). I refactored it a bit (decided it was better to have separate helper functions, instead of local lambda helpers), but you can see in fix_up it's very much procedural \$\endgroup\$ – Raekye Jan 23 '14 at 21:27
3
\$\begingroup\$

The code looks much better after refactoring!

Additionally I'd strongly suggest to keep line lengths within some limit. Usual choices are something between 72 and 80. There are two reasons for it:

  • People with smaller screens aren't able to see a piece code at once and have to scroll. This makes reading the code next to impossible (like when your snippet here on SO doesn't fit into its frame).
  • Even for a person with a wide screen it's difficult to read text with long lines. It's hard for eyes to focus where the next line starts.

Don't be afraid to use short identifiers, if theirs scope is limited to a short function. For example, in my opinion this

rotate_left :: Tree a -> Tree a
rotate_left (Branch l v (Branch rl rv rr _rc) c)
          = (Branch (Branch l v rl Red) rv rr c)

is more readable, and it's visually easy to spot what is going on.

One thing that will also help you make the code shorter more readable is using as-patterns (see also this question):

do_add :: (Ord a) => Tree a -> a -> Tree a
do_add branch@(Branch left node right colour) val
    | val < node = (Branch (add left val) node right colour)
    | val > node = (Branch left node (add right val) colour)
    | otherwise = branch -- HERE: we don't need to recreate the node
do_add Leaf val = (Branch Leaf val Leaf Black)

It also can give a small performance boost as we don't re-create objects identical to those we pattern match on.

Concerning fix_up: Let's try to factor out common and duplicate code. The common pattern is that we check some conditions on the sub-nodes of a node and if it's true, we apply a function or it. Otherwise we keep it intact. We can split this idea into two functions - one that is general, and other that is then specialized for branches:

fix_up :: Tree a -> Tree a
fix_up =
    onBr [is_red]                         [is_red] flip_colours .
    onBr [is_red, is_red . get_left_node] []       rotate_right .
    onBr [not . is_red]                   [is_red] rotate_left
  where
    -- Apply a function on a branch, if its left and right subnodes match
    -- given predicates.
    onBr :: [Tree a -> Bool] -> [Tree a -> Bool]
         -> (Tree a -> Tree a) -> (Tree a -> Tree a)
    onBr lps rps = on (\b -> all ($ get_left_node b) lps
                          && all ($ get_right_node b) rps)
    -- Apply a function on a value, if it matches a predicate.
    on :: (a -> Bool) -> (a -> a) -> (a -> a)
    on p f x | p x       = f x
             | otherwise = x

This allows us to represent the whole operation as the composition of several functions. And it localizes bindings of subnodes (l and r) to the predicates, which again helps readability.


I'd also like to draw your attention to AA trees. They have very similar performance as red-black trees, but are simpler and easier to implement.

\$\endgroup\$
  • \$\begingroup\$ Thanks for all the recommendations! I remember reading about short variable names in functional languages from Scala but never made the connection to do it here. I was also looking for the as pattern (which you can do with the bind operator in Erlang) but didn't know what it was called. The predicates are cool too. Finally thanks for pointing me to AA trees! \$\endgroup\$ – Raekye Jan 27 '14 at 0:16
  • \$\begingroup\$ PS: I just understood your use of the $... wow \$\endgroup\$ – Raekye Jan 27 '14 at 22:20
  • \$\begingroup\$ @Raekye I forgot to explain it. ($ x) is a function that takes another function and applies it on x. It's a partial application of ($), supplying its second argument. \$\endgroup\$ – Petr Pudlák Jan 27 '14 at 22:26
  • \$\begingroup\$ Haha yeah no worries I enjoyed figuring it out myself; really felt like a divine revelation :P \$\endgroup\$ – Raekye Jan 27 '14 at 22:33
  • \$\begingroup\$ @Raekye Interestingly, ($) is just a specialization of id. It'd be possible to also write (id x), although that'd be quite uncommon and less readable. \$\endgroup\$ – Petr Pudlák Jan 27 '14 at 22:35

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.