Treap with implicit keys

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
{-# LANGUAGE BangPatterns #-}

import Data.List
import Data.Function (on)
import Data.Ord (comparing)
import Data.Int (Int64)

priorities = map (\v -> v * 1000000000000000 `mod` (maxBound :: Int))
-- 'random' priorities

createTree a = foldl' f Null $ zip a (priorities a)
    where f s (v, p) = merge s node
                        where node = Tree { left = Null
                                          , right = Null
                                          , priority = p
                                          , size = 1
                                          , value = fromIntegral v
                                          , minim = fromIntegral v
                                          , inc = 0
                                          }

n = 200000
m = 200000

main = do let a = [1..n] :: [Int]
              t = createTree a
              s = foldr (.) id (replicate m splitAndMerge) t
          print (getSize t)
          print (height t)
          print (getSize s)
          print (height s)

splitAndMerge t = merge t1 t2
    where (t1, t2) = split (div n 2) t

data Tree = Null |
            Tree { left     :: {-# UNPACK #-} !Tree
                 , right    :: {-# UNPACK #-} !Tree
                 , priority :: {-# UNPACK #-} !Int
                 , size     :: {-# UNPACK #-} !Int
                 , value    :: {-# UNPACK #-} !Int64
                 , minim    :: {-# UNPACK #-} !Int64
                 , inc      :: {-# UNPACK #-} !Int64
                 }

getSize Null = 0
getSize t = size t

getMin Null = maxBound :: Int64
getMin t = minim t + inc t

height Null = 0
height Tree {left = l, right = r} = 1 + max (height l) (height r)

instance Eq Tree where
    (==) = (==) `on` priority
instance Ord Tree where
    compare = comparing priority

push t@Tree {inc = 0} = t
-- In a real program there will be another binding for push

update :: Tree ->  Tree
update t@Tree {left = l, right = r, value = v} =
       t { minim = min v $ min (getMin l) (getMin r),
           size  = 1 + getSize l + getSize r }

merge Null t = t
merge t Null = t
merge first' second'
    | first < second = let merged = merge (right first) second
                       in update $ first { right = merged }
    | otherwise      = let merged = merge first (left second)
                       in update $ second {left = merged}
    where !first = push first'
          !second = push second'

split :: Int -> Tree -> (Tree, Tree)
split 0 t = (Null, t)
split count t'
    | count <= getSize (left t)  = let (!t1, !t2) = split count $ left t
                                       !t3 = update $ t {left = t2}
                                   in (t1, t3)

    | otherwise = let (!t2, !t3) = split (count - 1 - getSize (left t)) (right t)
                      !t1 = update $ t {right = t2}
                  in (t1, t3)
    where !t = push t'