Competitive programming in Haskell: sorting tree shapes

In my previous post I challenged you to solve this problem, which essentially asks how many distinct binary tree shapes are created when we take lists of numbers and build a tree from each by repeated binary search tree insertion.

Incidentally, this problem was from the 2016 ICPC world finals (probably one of the easiest ICPC world finals problems ever!).

Several commenters solved it, and all with essentially the same solution. First we need to build some binary search trees by repeated insertion, which we can do by creating a binary tree type and insertion function and then doing some left folds over the input lists. Next, we need to classify the resulting trees by their shape. One obvious method would be to write a function which compares two binary trees to see if they have the same shape; use nubBy to remove duplicates; then count how many trees remain. This would take O(n^2) time, but since there are only at most 50 trees, with at most 20 values in each, this should easily fit within the very geneous time limit of 5 seconds. (This is an understatement; my implementation of this approach runs in 0.01s.)

However, there’s a different solution which is both asymptotically faster and less code! The key idea is that if we make the Tree type polymorphic, and an instance of Functor (by writing our own instance, or, even better, using the DeriveFunctor extension), then after building the trees we can turn them into literal tree shapes by replacing the values they contain with (). Moreover, since GHC can also derive an Ord instance for our Tree type, we can then count the distinct tree shapes in O(n \lg n) time, either using a combination of sort, group, and length, or by throwing them all into a Set and asking for its size.

Here’s my solution:

{-# LANGUAGE DeriveFunctor #-}

import Control.Arrow
import Data.List

main = interact $
  lines >>> drop 1 >>> map (words >>> map read) >>> solve >>> show

solve :: [[Int]] -> Int
solve = map (foldl' (flip ins) Empty >>> (() <$)) >>> sort >>> group >>> length
                                           -- or: >>> S.fromList >>> S.size

data Tree a = Empty | Node a (Tree a) (Tree a)
  deriving (Show, Eq, Ord, Functor)

ins :: Ord a => a -> Tree a -> Tree a
ins a Empty = Node a Empty Empty
ins a (Node x l r)
  | a < x     = Node x (ins a l) r
  | otherwise = Node x l (ins a r)

Honestly I’m not sure what the nicest way to solve this problem in something like Java or C++ would be. In Java, I suppose we would have to make a class for trees, implement equals and compareTo methods which compare trees by shape, and then put all the trees in a TreeSet; or else we could implement hashCode instead of compareTo and use a HashSet. The thing that makes the Haskell solution so much nicer is that the compiler writes some of the code for us, in the form of derived Functor and Ord instances.

For Friday, I invite you to solve Subway Tree System, a nifty problem which is more difficult but has some similar features!

About Brent

Assistant Professor of Computer Science at Hendrix College. Functional programmer, mathematician, teacher, pianist, follower of Jesus.
This entry was posted in competitive programming, haskell and tagged , , , , , , , . Bookmark the permalink.

18 Responses to Competitive programming in Haskell: sorting tree shapes

  1. blaisepascal2014 says:

    The two important lines in my solution were:

    data Subway = Subway [Subway] deriving (Show, Ord, Eq)
    canonical (Subway xs) = Subway $ sort (map canonical xs)

    Everything else was wrangling data in and out.

    • Brent says:

      Nice! As you’ll see in my post tomorrow, I chose a data representation that did the canonicalization for me automatically. But actually I think yours is simpler!

      • blaisepascal2014 says:

        Thanks. I don’t know exactly what the canonical form is, as I don’t know what the derived order of Subway is.

        Actually, the third important line was:

        parseSubway = Subway (string “0” *> many parseSubway <* string "1")

        which *almost* parses the input strings properly (I have to wrap the strings in a "0" and "1" to get it to work).

    • Anonymous says:

      Thanks for the hint! I implemented my own `Ord` instance, and used a similar approach to implement my own `Eq` instance, but something didn’t work and my program failed the second trial on kattis.

      With your hint my program passed!

  2. Justin says:

    These are fun!I spent most of my time figuring out how to build the tree from the traversal given. I started with a continuation-passing style solution, but as usual that hurt my brain a lot, so I switched to using a Zipper while folding over the list of steps. That worked out quite nicely!

    data Zipper a = Zipper { focus :: a, up :: Maybe (Zipper a) }
    data Tree a = Tree a [Tree a]
    deriving (Show, Ord, Eq)

    insert :: [Int] -> Tree Int
    insert xs =
    let (Zipper result _) = foldl go (Zipper (Tree (-1) []) Nothing) xs
    in result
    go :: Zipper (Tree Int) -> Int -> Zipper (Tree Int)
    go (Zipper focus (Just (Zipper (Tree a leaves) up))) 1 = Zipper (Tree a (focus : leaves)) up
    go zipper@(Zipper focus Nothing) 1 = zipper
    go zipper@(Zipper focus@(Tree d _) up) _ = Zipper (Tree (d – 1) []) (Just zipper)

  3. Globules says:

    It pains me to look at it, but it passed Kattis…

    import Data.List (foldl', sort)
    import Data.Tree (Tree(..), drawTree, levels)
    mkTree :: [Char] -> Tree Int
    mkTree = head . foldl' step [Node 0 []]
      where step ns '0' = Node 0 [] : ns
            step (n@(Node i _) : Node i' cs' : ns) '1' = Node (1+i+i') (n : cs') : ns
    -- ugh
    canon :: Tree Int -> [[Int]]
    canon = sort . map sort . levels
    same :: Tree Int -> Tree Int -> Bool
    same t0 t1 = canon t0 == canon t1
  4. Globules says:

    I’m dubious about the canon function in my solution. I suspect a better approach might be to define my own Tree type whose nodes store their children in some ordered container, then “same” becomes just an equality check. But, it’s bedtime…

  5. Anonymous says:

    My solution :
    Don’t like it much, there has to be a better way.
    Thanks a lot Brent, this is fun!

  6. Arnaud Bailly says:

    My solution: Not very elegant I am afraid, I was stuck in the tree reconstruction from the depth-first traversal.

  7. Pingback: Competitive programming in Haskell: building unordered trees | blog :: Brent -> [String]

  8. My solution is more pedagogical. It uses zippers and a hylomorphism. source.

  9. shaurya gupta says:

    Ooo this is exactly how I solved it! I suddenly don’t feel like a Haskell outsider.

      • shaurya gupta says:

        In C++, I would have either grouped them using a simple “hasSameStructure” function (for smaller constraints) or used a hashing function to hash the trees (larger constraints).

        • Brent says:

          Ah, a hashing function is a nice idea. So you mean you would specifically combine hashes of subtrees using a commutative and associative operation (e.g. adding) so you would get the same hash no matter their order?

Leave a Reply

Fill in your details below or click an icon to log in: Logo

You are commenting using your account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

This site uses Akismet to reduce spam. Learn how your comment data is processed.