/
redblack.hs
106 lines (88 loc) · 3.1 KB
/
redblack.hs
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
{-# LANGUAGE GADTs #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
import Data.Unfoldable
import Data.Unfolder
import Data.Maybe
import System.Random
import Data.List (intercalate)
-- Red-Black tree implementation adapted from https://gist.github.com/2660297
data Nat = Zero | Succ Nat
data NatW :: Nat -> * where
ZeroW :: NatW 'Zero
SuccW :: NatW n -> NatW ('Succ n)
data RedBlack = Black | Red
data RedBlackTree a where
T :: Node 'Black n a -> RedBlackTree a
deriving instance Show a => Show (RedBlackTree a)
data Node :: RedBlack -> Nat -> * -> * where
Leaf :: Node 'Black 'Zero a
B :: Node cL n a -> a -> Node cR n a -> Node 'Black ('Succ n) a
R :: Node 'Black n a -> a -> Node 'Black n a -> Node 'Red n a
deriving instance Show a => Show (Node c n a)
instance Unfoldable RedBlackTree where
unfold = u ZeroW
where
u :: forall n f a. (Unfoldable (Node 'Black n), Unfolder f)
=> NatW n -> f a -> f (RedBlackTree a)
u n fa = choose
[ T <$> (unfold :: f a -> f (Node 'Black n a)) fa
, u (SuccW n) fa
]
instance Unfoldable (Node 'Black 'Zero) where
unfold _ = choose [ pure Leaf ]
instance Unfoldable (Node 'Black n) => Unfoldable (Node 'Black ('Succ n)) where
unfold = u
where
u :: forall f a. Unfolder f => f a -> f (Node 'Black ('Succ n) a)
u fa = choose
[ B <$> b <*> fa <*> b
, B <$> b <*> fa <*> r
, B <$> r <*> fa <*> b
, B <$> r <*> fa <*> r
]
where
r :: f (Node 'Red n a)
r = unfold fa
b :: f (Node 'Black n a)
b = unfold fa
instance Unfoldable (Node 'Black n) => Unfoldable (Node 'Red n) where
unfold fa = choose [ R <$> unfold fa <*> fa <*> unfold fa ]
rbtree :: Int -> RedBlackTree Int
rbtree l = fromJust $ fromList [0..l]
rbtreeShapes :: Int -> [RedBlackTree ()]
rbtreeShapes d = limitDepth d unfold_
randomRBTree :: IO (RedBlackTree Bool)
randomRBTree = getStdRandom randomDefault
newtype Formula = Formula [Integer]
instance Show Formula where
show (Formula xs)
| all (== 0) xs = "0"
| otherwise = intercalate " + " $ map showOne $ filter ((/=0) . fst) $ zip xs [(0::Int)..]
where
showOne (x, 0) = show x
showOne (1, 1) = "x"
showOne (x, 1) = show x ++ "x"
showOne (1, n) = "x^" ++ show n
showOne (x, n) = show x ++ "x^" ++ show n
add, mul :: [Integer] -> [Integer] -> [Integer]
add xs [] = xs
add [] ys = ys
add (x:xs) (y:ys) = (x + y) : add xs ys
mul _ [] = []
mul [] _ = []
mul (x:xs) (y:ys) = (x * y) : add (map (x*) ys) (add (map (y*) xs) (0 : mul xs ys))
instance Num Formula where
fromInteger x = Formula [x]
Formula xs + Formula ys = Formula (add xs ys)
Formula xs * Formula ys = Formula (mul xs ys)
varX :: Formula
varX = Formula [0, 1]
-- See http://oeis.org/A001137
rbFormula :: Int -> NumConst Formula (RedBlackTree a)
rbFormula d = ala (limitDepth d) unfold (NumConst varX)