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!