Submission #286814

---

Source Code Expand

{-# LANGUAGE ScopedTypeVariables #-}

import           Control.Applicative
import           Control.Monad
import           Data.Vector         ((!))
import qualified Data.Vector         as V

main :: IO ()
main = do
  ([n, x] :: [Int]) <- map read . words <$> getLine
  (h :: [Int]) <- map read . words <$> getLine

  (es :: [[Int]]) <- replicateM (n - 1) $ map read . words <$> getLine

  let g :: V.Vector [Int]
      g = V.accum (flip (:)) (V.replicate n [])
          $ concat [ [(a-1, b-1), (b-1, a-1)] | [a, b] <- es ]

      hv :: V.Vector Int
      hv = V.fromList h

      solve cur par
        = sum
          [ solve c cur + 2
          | c <- g!cur
          , c /= par
          , cnt c cur > 0
          ]

      cnt cur par
        = (hv!cur) +
          sum
          [ cnt c cur
          | c <- g!cur
          , c /= par
          ]

  print $ solve (x-1) (-1)

  return ()

Submission Info

Judge Result