Data.Tree には木も森もある
競プロ典型 90 問 39 日目 - Tree Distance(★5)
import Control.Monad
import qualified Data.ByteString.Char8 as C
import Data.Graph
import Data.List
import Data.Tree
main = readLn >>= \n -> replicateM (n-1) ints >>= print . sol n
ints = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine
sol n es = sum . tmap ((*) <*> (-) n) $ go t
where
g = buildG (1,n) $ concatMap (\[u,v] -> [(u,v),(v,u)]) es
t = head $ dff g
tmap f (Node x ts) = Node (f x) (map (tmap f) ts)
go (Node _ []) = Node 1 []
go (Node x ts) = Node x' ts'
where
ts' = map go ts
x' = succ . sum $ map rootLabel ts'
いいなと思ったら応援しよう!
ありがとう