How to implement an AVL tree efficient enough in Haskell?

Question

upd: I am very sorry about my mistake. The old version of data is download from luogu, but the website do NOT allow me to download the data of the worst case. In fact, when facing same amount of operations, the memory usage depends on the type of operations and their orders. I re-generate some data, and update the link and profiling logs.

---

I am new in Haskell, and I try to implement an AVL tree. To test my code, I submit it to an Online Judge, which has a problem of balanced-tree. However, I get a MLE (Memory Limit Exceeded).

In fact, when facing 10^5 times of operations, my code requires 58 MB memories, while native C++ code only need about 2 to 3 MB. The speed of C++ code is also 10x faster.

I wonder if haskell can deal with such large data structure while having performance close to native code. What should I do?

Here is my code:

type Info = (Int, Int)
data Tree a = Empty | Ord a => Node (Tree a) a Info (Tree a)

just :: Maybe a -> a
just (Just a) = a

height :: Tree a -> Int
height Empty = 0
height (Node _ _ (h,_) _) = h

size :: Tree a -> Int
size Empty = 0
size (Node _ _ (_,s) _) = s

maintainInfo :: Tree a -> Tree a
maintainInfo (Node ls v _ rs) = Node ls v (max (height ls) (height rs) + 1, size ls + size rs + 1) rs

rotateL :: Tree a -> Tree a
rotateL (Node (Node ll lv _ lr) v _ rs) = 
    maintainInfo (Node ll lv (0,0) (maintainInfo (Node lr v (0,0) rs)))

rotateR :: Tree a -> Tree a
rotateR (Node ls v _ (Node rl rv _ rr)) = 
    maintainInfo (Node (maintainInfo (Node ls v (0,0) rl)) rv (0,0) rr)

rotateL_ :: Tree a -> Tree a
rotateL_ (Node (Node ll lv _ lr) v _ rs)
    | height ll >= height lr  = rotateL (Node (Node ll lv (0,0) lr) v (0,0) rs)
    | otherwise               = rotateL (Node (rotateR (Node ll lv (0,0) lr)) v (0,0) rs)

rotateR_ :: Tree a -> Tree a
rotateR_ (Node ls v _ (Node rl rv _ rr))
    | height rl > height rr   = rotateR (Node ls v (0,0) (rotateL (Node rl rv (0,0) rr)))
    | otherwise               = rotateR (Node ls v (0,0) (Node rl rv (0,0) rr))

maintain :: Tree a -> Tree a
maintain (Node ls v info rs)
    | height ls > height rs + 1 = rotateL_ (Node ls v info rs)
    | height rs > height ls + 1 = rotateR_ (Node ls v info rs)
    | otherwise                 = maintainInfo (Node ls v info rs)

insert :: Ord a => a -> Tree a -> Tree a
insert val Empty = Node Empty val (1,1) Empty
insert val (Node ls v _ rs)
    | val < v     = maintain (Node (insert val ls) v (0,0) rs)
    | val >= v    = maintain (Node ls v (0,0) (insert val rs))

iter :: Tree a -> [a]
iter Empty = []
iter (Node ls v _ rs) = iter ls ++ v : iter rs

maxi :: Tree a -> Maybe a
maxi Empty                 = Nothing
maxi (Node _ v _ Empty)    = Just v
maxi (Node _ _ _ rs)     = maxi rs

mini :: Tree a -> Maybe a
mini Empty                 = Nothing
mini (Node Empty v _ _)    = Just v
mini (Node ls _ _ _)     = mini ls

eraseSwap :: Tree a -> Tree a
eraseSwap Empty = Empty
eraseSwap (Node ls _ _ Empty) = ls
eraseSwap (Node ls _ _ rs) = 
    let next = just (mini rs) in 
        maintain (Node ls next (0,0) (erase next rs))

erase :: a -> Tree a -> Tree a
erase _ Empty = Empty
erase val (Node ls v _ rs)
    | val < v   = maintain (Node (erase val ls) v (0,0) rs)
    | val > v   = maintain (Node ls v (0,0) (erase val rs))
    | val == v  = eraseSwap (Node ls v (0,0) rs)

next :: a -> Tree a -> Maybe a
next _ Empty = Nothing
next val (Node ls v _ rs) 
    | val < v   = Just (maybe v (min v) (next val ls))
    | val >= v  = next val rs

prev :: a -> Tree a -> Maybe a
prev _ Empty = Nothing
prev val (Node ls v _ rs)
    | val <= v  = prev val ls
    | val > v   = Just (maybe v (max v) (prev val rs))

rank :: a -> Tree a -> Int
rank _ Empty = 1
rank val (Node ls v _ rs)
    | val <= v  = rank val ls
    | val > v   = size ls + 1 + rank val rs

select :: Int -> Tree a -> Maybe a
select _ Empty = Nothing
select k (Node ls v _ rs)
    | size ls >= k      = select k ls
    | size ls == k - 1  = Just v
    | otherwise         = select (k - size ls - 1) rs

update :: [Int] -> Tree Int -> Tree Int
update [1,x] = insert x 
update [2,x] = erase x

ask :: [Int] -> Tree Int -> Int 
ask [3,x] tree = rank x tree
ask [4,x] tree = just (select x tree)
ask [5,x] tree = just (prev x tree)
ask [6,x] tree = just (next x tree)

solve :: Int -> Tree Int -> IO ()
solve 0 _ = return ()
solve n tree = do
    inputQuery <- getLine
    let l = map read (words inputQuery) :: [Int]
    if head l < 3 then do 
        solve (n-1) (update l tree)
    else do
        print (ask l tree)
        solve (n-1) tree
    return ()

main :: IO ()
main = do 
    inputN <- getLine
    let n = read inputN :: Int
    solve n Empty
    return ()

I humbly accept all suggestions, including those that are not aimed at performance but at coding style.

---

To test my code, I submit it to Luogu P3369, which is in Chinese. The simple translation is as follow:

You need to dynamically maintain a re-set M and provide the following operations:

1. Insert a number x into M.
2. Delete a number x from M (if there are multiple identical numbers, only delete one).
3. Query how many numbers in M are smaller than x, and add one to the answer.
4. Query the number that ranks in the x-th position after M is sorted from small to large.
5. Query the predecessor of x in M (the predecessor is defined as the largest number that is smaller than x).
6. Query the successor of x in M (the successor is defined as the smallest number that is greater than x).

• For operations 3, 5, and 6, it is not guaranteed that the number x exists in the current re-set.
• The first line is n, which indicates the number of operations. Each of the following n lines has two numbers opt and x. opt indicates the sequence number of the operation (1 <= opt <= 6)
• For operations 3,4,5,6, each line outputs a number, indicating the corresponding answer.

If anyone has difficulty to login luogu, here is also some data files to download:

• inputs
• outputs

Here is the output of ghc -O2 -rtsopts main.hs && ./main < avl.in > tmp.out +RTS -s:

   1,579,041,840 bytes allocated in the heap
     309,624,880 bytes copied during GC
      19,305,064 bytes maximum residency (21 sample(s))
         186,640 bytes maximum slop
              58 MiB total memory in use (0 MB lost due to fragmentation)

                                     Tot time (elapsed)  Avg pause  Max pause
  Gen  0       360 colls,     0 par    0.324s   0.329s     0.0009s    0.0055s
  Gen  1        21 colls,     0 par    0.419s   0.425s     0.0202s    0.0500s

  INIT    time    0.001s  (  0.001s elapsed)
  MUT     time    0.897s  (  0.907s elapsed)
  GC      time    0.743s  (  0.754s elapsed)
  EXIT    time    0.000s  (  0.009s elapsed)
  Total   time    1.641s  (  1.671s elapsed)

  %GC     time       0.0%  (0.0% elapsed)

  Alloc rate    1,759,782,364 bytes per MUT second

  Productivity  54.7% of total user, 54.3% of total elapsed

Answer 1

As I said in a comment, I can't verify that your program as written uses 40MB - your profiling data and my own runs of your program show 3MB, in the same range as the C++ performance you refer to. Even with -O0, I see only 6MB residency for your input file with 10k lines.

So I can't make any recommendations that are guaranteed to greatly reduce memory usage: it looks reduced enough to me already. But the first few things I would try are:

Improve strictness

1. Add {-# LANGUAGE Strict #-} to the top of the file. You don't need any laziness, and this will help prevent space leaks. When I try this I see no change in memory usage, and only a 5% or so speedup (not really statistically significant, I imagine)
2. Replace your Info tuple with a real data type. This will let the runtime avoid using the existing tuple type, which is of course lazy, and use a type that's equivalent but can be strict (thanks to the Strict language extension).
3. Replace all of your (0,0) expression with a reference to a top-level constant defined as untracked = Info 0 0. I think you did this in the version of this question you posted to Stack Overflow, calling it nilInfo instead. It's a good idea, so you don't waste time building this same boring value over and over. When I combine these 3 improvements, I again see the same memory usage, and a larger but still not significant speedup.

Reduce use of dummy Info values

Next, I note that many of your uses of this untracked "dummy Info" value are followed immediately by a call to maintainInfo, which ignores the Info field and constructs its own. So you are building a Tree node just to tear it down and build a new one. Instead, define a function that takes all the other Tree fields as parameters, and constructs a Node with the appropriate Info:

mkNode :: Ord a => Tree a -> a -> Tree a -> Tree a
mkNode ls v rs = Node ls v (Info (max (height ls) (height rs) + 1)
                                 (size ls + size rs + 1))
                      rs

The call sites are easy to update: just replace maintainInfo (Node x y untracked z) with mkNode x y z.

You should be able to do something similar with the rotate family of functions, and maintain, all of which ignore Info on the nodes they are given, and construct new Info values based on their subtrees after rotation, but still waste time constructing Info metadata to pass around to each other. That's more surgery than I'm interested in doing myself for this answer, though.

Participate in Prelude typeclasses

Instead of bespoke functions iter, maxi, and mini, I suggest participating in Foldable. Its maximum and minimum functions don't match signature with your maxi and mini - you return a Maybe, in the modern style, while these functions with much older specifications simply crash when given an empty list. Fortunately, you never call maxi at all, and you assume that mini always returns a Just anyway. So you might as well make your Tree Foldable. This won't help performance at all, but it's just a good practice.

extremum :: String -> (forall a. Tree a -> Tree a) -> (forall a. Ord a => Tree a -> a)
extremum name next = impl
  where impl Empty = error (name <> ": empty Tree")
        impl n = go n
        go n = case next n of
          Empty -> let Node _ v _ _ = n
                   in v
          n' -> go n'

instance Foldable Tree where
  foldMap _ Empty = mempty
  foldMap f (Node ls v _ rs) = foldMap f ls <> f v <> foldMap f rs
  maximum = extremum "maximum" $ \(Node _ _ _ rs) -> rs
  minimum = extremum "minimum" $ \(Node ls _ _ _) -> ls
  length = size

I am being more "clever" than necessary with the extremum type signature and implementation. If you don't understand it or don't like it, a simpler way to implement maximum and minimum would be fine of course - it's nothing critical.

More idiomatic use of comparisons in guard clauses

In erase, you compare two values three times, checking for <, then for >, then for ==. Instead, use compare, which returns a type with 3 values, so you can simply fall into the right case directly.

erase :: a -> Tree a -> Tree a
erase _ Empty = Empty
erase val (Node ls v _ rs) = case compare val v of
  LT -> maintain (Node (erase val ls) v untracked rs)
  GT -> maintain (Node ls v untracked (erase val rs))
  EQ -> eraseSwap (Node ls v untracked rs)

In many other functions, you compare with two mutually exclusive conditions, such as >= and then <. That could be replaced with compare in a similar way (match LT in the first case and _ for the other case), but it is more common to use otherwise instead. Either way, you remove a redundant comparison and make the code easier to understand.

next :: a -> Tree a -> Maybe a
next _ Empty = Nothing
next val (Node ls v _ rs)
    | val < v   = Just (maybe v (min v) (next val ls))
    | otherwise = next val rs

Your select implementation can avoid some redundant function calls by using a where clause to span all the guards:

select :: Int -> Tree a -> Maybe a
select _ Empty = Nothing
select k (Node ls v _ rs)
    | lSize >= k      = select k ls
    | lSize == k - 1  = Just v
    | otherwise       = select (k - lSize - 1) rs
  where lSize = size ls

Consider a different data structure

AVL trees are a fine data structure to implement for an exercise like this. But if you are worried about performance, they are not famous for being the fastest tree to keep balanced. Their invariant is very strict, and so a lot of time is spent keeping them perfectly balanced, while a data structure with a more relaxed invariant can let the trees get a little lopsided without letting them get too unbalanced. Red/black trees are another classic, simple example, and will spend less time maintaining balance at the cost of a slightly deeper average tree.

Or you could use some types from the standard library. Data.IntMap should support all the features you need, for example, and I imagine it's more efficient than anything you could do yourself. Or there's probably a fun way to build this yourself out of a FingerTree, which would be somewhere in between doing it all yourself and having it all done for you. It depends on how much you want to do by hand for the exercise, and how much you're happy to have handled for you.