A Painting Puzzle

graph.dot

digraph G {
  size = "4,4";
  "[1,1,1,1]" -> "[1,1,2]" [label="1"];
  "[1,1,2]"   -> "[1,1,2]" [label="1/2"];
  "[1,1,2]"   -> "[1,3]"   [label="1/3"];
  "[1,1,2]"   -> "[2,2]"   [label="1/6"];
  "[1,3]"     -> "[1,3]"   [label="1/2"];
  "[1,3]"     -> "[2,2]"   [label="1/4"];
  "[1,3]"     -> "[4]"     [label="1/4"];
  "[2,2]"     -> "[1,3]"   [label="2/3"];
  "[2,2]"     -> "[2,2]"   [label="1/3"];
  "[4]"       -> "[4]"     [label="1"];
}

painting.hs

import Control.Arrow          (second, (&&&))
import Data.Either            (rights)
import Data.List              (sort, group, permutations, delete)
import Data.Map.Strict as Map (Map, empty, fromList, keys, elems, insert, 
                               findWithDefault)
import Data.Matrix as Matrix  (Matrix, fromList, multStd, inverse, identity, submatrix, nrows)
import Data.Ratio             (Ratio, (%))
import Data.Set as Set        (Set, toList, difference, fromList)
import Prelude as P
import System.Environment     (getArgs)

pDist :: [Int] -> Map [Int] (Ratio Int)
pDist = freqMap . P.map decorate . allPairings . toUniqueRepresentation
  where toUniqueRepresentation      :: [Int] -> [Int]
        toUniqueRepresentation node = concat $ zipWith replicate node [1..]

        allPairings    :: [Int] -> [[Int]]
        allPairings ns = [ x : y : delete y (delete x ns) 
                         | x <- ns, y <- delete x ns
                         ]

        decorate       :: [Int] -> [Int]
        decorate       = sort . P.map length . group . sort . paint

        paint          :: [Int] -> [Int]
        paint (x:y:xs) = x:x:xs

        freqMap    :: [[Int]] -> Map [Int] (Ratio Int)
        freqMap xs = Map.fromList
                   $ P.map ( second (% length xs) . (head &&& length) )
                   $ group $ sort xs

makeArrows      :: [Int] -> Map [Int] (Map [Int] (Ratio Int))
makeArrows seed = until finished step $ Map.insert seed (pDist seed) Map.empty
  where finished    :: Map [Int] (Map [Int] (Ratio Int)) -> Bool
        finished    = P.null . diff

        step        :: Map [Int] (Map [Int] (Ratio Int)) -> Map [Int] (Map [Int] (Ratio Int))
        step arrows = Map.insert node (pDist node) arrows
          where node :: [Int]
                node = head $ Set.toList $ diff arrows

        diff        :: Map [Int] (Map [Int] (Ratio Int)) -> Set [Int]
        diff arrows = Set.difference (js arrows) (is arrows)

        is          :: Map [Int] (Map [Int] (Ratio Int)) -> Set [Int]
        is          = Set.fromList . Map.keys

        js          :: Map [Int] (Map [Int] (Ratio Int)) -> Set [Int]
        js          = Set.fromList . concatMap Map.keys . Map.elems


makeMatrix        :: Map [Int] (Map [Int] (Ratio Int)) -> Matrix (Ratio Int)
makeMatrix arrows = Matrix.fromList n n list
  where list   :: [Ratio Int]
        list   = [ Map.findWithDefault (0%1) j $ Map.findWithDefault Map.empty i arrows 
                 | i <- nodes, j <- nodes
                 ]
        n      :: Int
        n      = length nodes
        nodes :: [[Int]]
        nodes = Map.keys arrows

expectedValue   :: Matrix (Ratio Int) -> Matrix (Ratio Int)
expectedValue p = P.foldr1 multStd [tau, inv, one]
  where tau = Matrix.fromList 1 n (1:[0,0..]) -- {1,0,0..}
        one = Matrix.fromList n 1 [1,1..]     -- {{1},{1},...}
        inv = head $ rights $ (:[]) 
            $ inverse (identity n - t)        -- (I - T)^(-1)
        t = submatrix 1 n 1 n p               -- P, but w/o last row / last col
        n = nrows p - 1                      

main = do
  args <- getArgs
  let seed = read (head args) :: [Int]
  print $ expectedValue $ makeMatrix $ makeArrows seed