Skip to content

Instantly share code, notes, and snippets.

@EncodePanda
Last active November 21, 2021 21:51
Show Gist options
  • Save EncodePanda/3a351c1783230c805c3bf733ae4ac21e to your computer and use it in GitHub Desktop.
Save EncodePanda/3a351c1783230c805c3bf733ae4ac21e to your computer and use it in GitHub Desktop.
First attempt to implement 2-3 tree as outlined in the "Finger Trees: A Simple General-purpose Data Structure" paper
-- λ> fromList [1]
-- Zero 1
-- λ> fromList [1..2]
-- Succ (Zero (Node2 1 2))
-- λ> fromList [1..4]
-- Succ (Succ (Zero (Node2 (Node2 1 2) (Node2 3 4))))
-- λ> fromList [1..8]
-- Succ (Succ (Succ (Zero (Node2 (Node2 (Node2 1 2) (Node2 3 4)) (Node2 (Node2 5 6) (Node2 7 8))))))
-- λ> head $ fromList [1..8]
-- 1
-- λ> tail $ fromList [1..8]
-- Just (Succ (Succ (Zero (Node3 (Node3 2 3 4) (Node2 5 6) (Node2 7 8)))))
-- λ> toList $ fromList [1..8]
-- [1,2,3,4,5,6,7,8]
module TwoThreeTree
where
import Prelude hiding (head, tail)
data Tree a =
Zero a
| Succ (Tree (Node a))
deriving Show
data Node a =
Node2 a a
| Node3 a a a
deriving Show
--------------------------------------------------------------------------------
-- | head is a safe function, as 'Tree a' represents a non-empty list
head :: Tree a -> a
head (Zero a) = a
head (Succ (Zero node)) = nodeHead node
head (Succ (Succ tree)) = nodeHead $ nodeHead $ head tree
nodeHead :: Node a -> a
nodeHead (Node2 a _) = a
nodeHead (Node3 a _ _) = a
--------------------------------------------------------------------------------
tail :: Tree a -> Maybe (Tree a)
tail (Zero _) = Nothing
tail (Succ (Zero (Node2 a b))) = Just $ Zero b
tail (Succ (Zero (Node3 a b c))) = Just $ Succ (Zero (Node2 b c))
tail (Succ (Succ (Zero node))) = case nodeTail node of
Left node -> Just $ Succ $ Zero node
Right node -> Just $ Succ $ Succ $ Zero node
tail (Succ (Succ tree)) = case nodeTail $ head tree of
Left node -> fmap (\t -> Succ $ node <: t) (Succ <$> tail tree)
Right node -> fmap (\t -> Succ $ Succ $ node <: t) (tail tree)
nodeTail :: (Node (Node a)) -> Either (Node a) (Node (Node a))
nodeTail (Node2 (Node2 a b) (Node2 c d)) = Left $ Node3 b c d
nodeTail (Node2 (Node2 a b) (Node3 c d e) ) = Right $ Node2 (Node2 b c) (Node2 d e)
nodeTail (Node2 (Node3 a b c) rest ) = Right $ Node2 (Node2 b c) rest
nodeTail (Node3 (Node2 a b) (Node2 c d) rest) = Right $ Node2 (Node3 b c d) rest
nodeTail (Node3 (Node2 a b) (Node3 c d e) rest) = Right $ Node3 (Node2 b c) (Node2 d e) rest
nodeTail (Node3 (Node3 a b c) rest1 rest2) = Right $ Node3 (Node2 b c) rest1 rest2
--------------------------------------------------------------------------------
infixr 5 <:
(<:) :: a -> Tree a -> Tree a
(<:) x (Zero a) = Succ (Zero $ Node2 x a)
(<:) x (Succ (Zero (Node2 a b))) = Succ (Zero $ Node3 x a b)
(<:) x (Succ (Zero (Node3 a b c))) = Succ (Succ $ Zero $ Node2 (Node2 x a) (Node2 b c))
(<:) x (Succ tree@(Succ _)) = case (head tree, tail tree) of
(Node2 a b, Nothing) -> Succ $ Zero $ Node3 x a b
(Node2 a b, Just tl) -> Succ $ Node3 x a b <: tl
(Node3 a b c, Nothing) -> Succ $ Succ $ Zero $ Node2 (Node2 x a) (Node2 b c)
(Node3 a b c, Just tl) -> Succ $ Node2 x a <: Node2 b c <: tl
--------------------------------------------------------------------------------
one a = Zero a
-- | fails from an empty list
fromList :: [a] -> Tree a
fromList [x] = Zero x
fromList (x:xs) = x <: fromList xs
toList :: Tree a -> [a]
toList tree = case tail tree of
Nothing -> [head tree]
Just t -> head tree : toList t
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment