12

31

[Haskell] 二分木が平衡かどうかを調べる関数

2012.12.31(12:51)

世界で闘うプログラミング力を鍛える150問
http://www.amazon.co.jp/dp/4839942390/
の問題4.1をHaskellで解いてみた。

問題 4.1 二分木が平衡かどうかを調べる関数を実装してください。 平衡木とは、どのノードの2つの部分木も、その高さの差が1以下である ような木であると定義します。

まず二分木を作り、それを使って問題を解く。

{-# OPTIONS -Wall -Werror #-}

module Tree where
import Test.HUnit
import Test.QuickCheck

-- | ----------------------- 木を植えよう ---------------------------
-- | すごいHaskell 7.7 再帰的なデータ構造
-- | Tree型の定義
data Tree a = Empty | Node a (Tree a) (Tree a)
            deriving Show

-- | 要素をTreeに追加する
-- | すでに存在していたら、そのNodeをそのまま返す
treeInsert :: (Ord a) => a -> Tree a -> Tree a
treeInsert x Empty = Node x Empty Empty
treeInsert x (Node a left right)
    | x == a    = (Node a left right)
    | x < a     = (Node a (treeInsert x left) right)
    | otherwise = (Node a left (treeInsert x right))

-- | 要素がTreeに含まれていればTrue
treeElem :: (Ord a) => a -> Tree a -> Bool
treeElem _ Empty = False
treeElem x (Node a left right)
    | x == a    = True
    | x < a     = treeElem x left
    | otherwise = treeElem x right

-- | ------------- 世界で闘うプログラミング力を鍛える150問 -----------------
-- | 問題 4.1 二分木が平衡かどうかを調べる関数を実装してください。
-- | 平衡木とは、どのノードの2つの部分木も、その高さの差が1以下である
-- | ような木であると定義します。
-- | ----------------------- アルゴリズム1 ---------------------------
-- | Treeの高さを求める
getHeight :: Tree a -> Int
getHeight Empty = 0
getHeight (Node _ left right)
    = max (getHeight left) (getHeight right) + 1

-- | Treeが平衡だったらTrue
isBalanced :: Tree a -> Bool
isBalanced Empty = True
isBalanced (Node _ left right)
    = if heightDiff > 1 then False
                        else isBalanced left && isBalanced right
         where heightDiff = abs $ getHeight left - getHeight right

-- | ----------------------- アルゴリズム2 ---------------------------
-- | leftとrightのどちらかが平衡でなければ-1を返す
-- | そうでなければleftとrightの高い方+1を返す
shortCut :: Int -> Int -> Int
shortCut leftHeight rightHeight
    | leftHeight  == (-1) = (-1)
    | rightHeight == (-1) = (-1)
    | (abs $ leftHeight - rightHeight) > 1 = (-1)
    | otherwise           = max leftHeight rightHeight + 1

-- | Treeの高さを求める
-- | ただし、leftとrightのどちらかが平衡でなければ-1を返す
checkHeight :: Tree a -> Int
checkHeight Empty = 0
checkHeight (Node _ left right) 
            = shortCut leftHeight rightHeight
                       where leftHeight = checkHeight left
                             rightHeight = checkHeight right

-- | Treeが平衡だったらTrue
isBalanced' :: Tree a -> Bool
isBalanced' t = checkHeight t /= (-1)

-- | ---------------------- テスト ---------------------------
-- | Treeを巡回してprintする
traverse :: (Show a) => Tree a -> IO ()
traverse Empty = putStr ""
traverse (Node a left right) = do traverse left
                                  putStrLn $ show a
                                  traverse right

-- | [a]を空のTreeに追加する
list2Tree :: (Ord a) => [a] -> Tree a
list2Tree nums = foldr treeInsert Empty nums

-- | Treeを巡回して[a]を作る
tree2List :: Tree a -> [a]
tree2List Empty = []
tree2List (Node a left right) = (tree2List left) ++ [a] ++ (tree2List right)

-- | ユニットテスト
balanceTest :: Test
balanceTest = test [
                (tree2List $ list2Tree ([]::[Int])) ~=? ([]::[Int]),
                (tree2List $ list2Tree ([2]::[Int])) ~=? ([2]::[Int]),
                (tree2List $ list2Tree ([1,2,3,4]::[Int])) ~=? ([1,2,3,4]::[Int]),
                (tree2List $ list2Tree ([1,3,4,2]::[Int])) ~=? ([1,2,3,4]::[Int]),
                (tree2List $ list2Tree ([1..100]::[Int])) ~=? ([1..100]::[Int]),
                (tree2List $ list2Tree ([100,99..1]::[Int])) ~=? ([1..100]::[Int]),
                (tree2List $ list2Tree "2013") ~=? ['0'..'3'],
                isBalanced myTree1 ~=? True,
                isBalanced myTree2 ~=? False,
                isBalanced' myTree1 ~=? True,
                isBalanced' myTree2 ~=? False,
                True ~=? True
    ] where myTree1 = list2Tree ([5,10,3,12,8]::[Int])  -- 型制約必要
            myTree2 = list2Tree ([4,3,1]::[Int])

-- | quickCheck
prop_balance :: [Int]->Bool
prop_balance nums = (isBalanced myTree == isBalanced' myTree)
                    where myTree = list2Tree nums

-- | main
main :: IO ()
main = do
       -- let nums = [5,10,3,12,8]
       -- let nums = [8,6,4,1,7,3,5] -- balanced
       let nums = ([8,7,6,5,4,3,1]::[Int]) -- un-balanced
       let t = foldr treeInsert Empty nums
       print t
       print $ treeElem 3 t
       print $ treeElem 2 t
       print $ getHeight t
       print $ isBalanced t
       print $ isBalanced' t
       traverse t
       mapM_ (putStrLn.show) $ scanr treeInsert Empty nums 
       putStrLn . tree2List $ list2Tree "Hello,World! 2013"
       _ <- runTestTT balanceTest
       quickCheck prop_balance
       return ()


プロフィール

島敏博

Shima Toshihiro 島敏博
信州アルプスハイランド在住。HaskellとElixirが好き。組み込みソフトウェアアーキテクト、C++プログラマ、山歩き、美術館巡り、和食食べ歩き、日本赤十字社救急法指導員、インデックス投資、クラシック音楽、SESSAME会員、状態マシン設計、モデル駆動開発、ソフトウェアプロダクトライン、Rubyist、実践ビジネス英語

■ ツイッター
http://twitter.com/saltheads
■ Facebook
http://www.facebook.com/saltheads
■ Qiita
http://qiita.com/saltheads

印刷する場合は、ブラウザの印刷メニューではなく、このページの上から3cmくらいの青いところにある、「印刷」を押してみてください。少しうまく印刷できます。まだ完全ではないのですが、これで勘弁してください。


カテゴリ
最新記事
月別アーカイブ
最新コメント
検索フォーム
リンク
sessame
RSSリンクの表示