Skip to content

Instantly share code, notes, and snippets.

@tanakh
Created June 1, 2014 16:17
Show Gist options
  • Save tanakh/ea1619a71bfea7e19ba5 to your computer and use it in GitHub Desktop.
Save tanakh/ea1619a71bfea7e19ba5 to your computer and use it in GitHub Desktop.
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
import Control.Applicative
import Control.Monad
import Data.Array.IO
import Data.Array.Unboxed
import qualified Data.ByteString.Char8 as C
import Data.Function (on)
import Data.List
import Data.Maybe (mapMaybe)
import Text.Printf (printf)
-- I/O
readInts :: C.ByteString -> [Int]
readInts = map fst . mapMaybe C.readInt . C.words
getInts :: IO [Int]
getInts = readInts <$> C.getLine
-- UnionFind
type UF = IOUArray Int Int
newUF :: Int -> IO UF
newUF n = newListArray (0, n-1) [0..n-1]
root :: UF -> Int -> IO Int
root uf v = do
p <- readArray uf v
when (p /= v) $ writeArray uf v =<< root uf p
readArray uf v
unite :: UF -> Int -> Int -> IO ()
unite uf v w = do
r <- root uf v
s <- root uf w
writeArray uf r s
----
main :: IO ()
main = do
[n, m] <- getInts
vs :: UArray Int Int <- listArray (0, n-1) <$> getInts
es <- replicateM m getInts
uf <- newUF n
sz :: IOUArray Int Int <- newArray (0, n) (1 :: Int)
let f !acc [v, u] = do
rv <- root uf v
ru <- root uf u
if rv == ru
then return acc
else do
unite uf v u
sv <- readArray sz rv
su <- readArray sz ru
rr <- root uf v
writeArray sz rr $ su + sv
return $ acc + fromIntegral (min (vs!v) (vs!u)) * fromIntegral su * fromIntegral sv
ans <- foldM f (0 :: Double)
$ sortBy (flip compare `on` \[v, u] -> min (vs!v) (vs!u))
$ map (map pred) es
printf "%.10f\n" (ans / fromIntegral n / fromIntegral (n-1) * 2 :: Double)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment