{- |
Description: Height-balanced binary trees with given number of nodes
Copyright: Copyright (C) 2021 Yoo Chung
License: GPL-3.0-or-later
Maintainer: dev@chungyc.org

Some solutions to "Problems.P60" of Ninety-Nine Haskell "Problems".
-}
module Solutions.P60 (heightBalancedTreesWithNodes) where

import           Problems.BinaryTrees
import           Problems.P59

-- | Construct all the height-balanced binary trees with a given number of nodes.
heightBalancedTreesWithNodes :: Int -> [Tree ()]
heightBalancedTreesWithNodes :: Int -> [Tree ()]
heightBalancedTreesWithNodes Int
n = (Int -> [Tree ()]) -> [Int] -> [Tree ()]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Int -> [Tree ()]
treesWithHeight [Int -> Int
minHeight Int
n..Int -> Int
maxHeight Int
n]
  where treesWithHeight :: Int -> [Tree ()]
treesWithHeight Int
h = (Tree () -> Bool) -> [Tree ()] -> [Tree ()]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) Int
n (Int -> Bool) -> (Tree () -> Int) -> Tree () -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree () -> Int
forall a. Tree a -> Int
treeSize) ([Tree ()] -> [Tree ()]) -> [Tree ()] -> [Tree ()]
forall a b. (a -> b) -> a -> b
$ Int -> [Tree ()]
heightBalancedTrees Int
h

-- | Minimum number of nodes in a height-balanced tree of height @h@.
minNodes :: Int -> Int
minNodes :: Int -> Int
minNodes Int
0 = Int
0
minNodes Int
1 = Int
1
minNodes Int
h = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
minNodes (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
minNodes (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2)

minHeight :: Int -> Int
minHeight :: Int -> Int
minHeight Int
n = Int -> Int
log2 (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

log2 :: Int -> Int
log2 :: Int -> Int
log2 Int
1 = Int
0
log2 Int
n = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
log2 (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

-- | Maximum height for a height-balanced tree with @n@ nodes.
maxHeight :: Int -> Int
maxHeight :: Int -> Int
maxHeight Int
n = Int -> Int -> Int
maxHeight' Int
n Int
0

maxHeight' :: Int -> Int -> Int
maxHeight' :: Int -> Int -> Int
maxHeight' Int
n Int
h
  | Int -> Int
minNodes Int
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Int
minNodes (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) = Int
h
  | Bool
otherwise                              = Int -> Int -> Int
maxHeight' Int
n (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)