hw7: add takeJ

[yorgey.git] / hw7 / JoinList.hs

diff --git a/hw7/JoinList.hs b/hw7/JoinList.hs
index b7a4aeb949599db702dfd0696d9e4cd2a4261581..7a353cded5f694c8e7fc606a1258bc8d8d86142a 100644 (file)
--- a/hw7/JoinList.hs
+++ b/hw7/JoinList.hs
@@ -26,18 +26,37 @@ indexJ :: (Sized b, Monoid b) =>
           Int -> JoinList b a -> Maybe a
 indexJ _ Empty = Nothing
 indexJ n _ | n < 0 = Nothing
-indexJ _ (Single _ a) = Just a
-indexJ n (Append _ l r) | n < (getSize (size (tag l))) = indexJ n l
+indexJ n (Single _ a) | n == 0 = Just a
+                      | otherwise = Nothing
+indexJ n (Append t l r) | n >= getSize (size t) = Nothing
+                        | n < 0 = Nothing
+                        | n < (getSize (size (tag l))) = indexJ n l
                         | otherwise = indexJ (n - getSize (size (tag l))) r
 
 -- 2. drop
 dropJ :: (Sized b, Monoid b) =>
          Int -> JoinList b a -> JoinList b a
+dropJ n x | n < 0 = x
 dropJ 0 x = x
 dropJ _ Empty = Empty
-dropJ _ (Single _ _) = Empty
-dropJ n x | (getSize (size (tag x))) < n = Empty
-dropJ n (Append _ l r) | n > (getSize (size (tag l))) =
+dropJ n v@(Single _ _) | n == 0 = v
+                       | otherwise = Empty
+dropJ n (Append t l r) | n > getSize (size t) = Empty
+                       | n > (getSize (size (tag l))) =
                            dropJ (n - getSize (size (tag l))) r
                        | otherwise = let lft = dropJ n l in
                                      Append ((tag lft) `mappend` (tag r)) lft r
+
+-- 3. take
+takeJ :: (Sized b, Monoid b) =>
+         Int -> JoinList b a -> JoinList b a
+takeJ 0 _ = Empty
+takeJ n _ | n < 0 = Empty
+takeJ _ Empty = Empty
+takeJ n v@(Single _ _) | n == 0 = v
+                       | otherwise = Empty
+takeJ n v@(Append t l r) | n > getSize (size t) = v
+                         | n < getSize (size (tag l)) = takeJ n l
+                         | otherwise = let rt = takeJ (n - getSize (size (tag l))) r in
+                                       Append ((tag l) `mappend` (tag rt)) l rt
+