見出し画像

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'

いいなと思ったら応援しよう!

karoyakani
ありがとう