Problem 217

217 Balanced Numbers
Problem 217 - Project Euler
探索範囲が広いので、全探索は無理。
また、個数ではなく、和が必要。
さて、どうDPを構成するか…
まず、半分の数字の和は高々、9*47で小さい。
和をキーとして考え、バランスナンバーの総和が必要だから、
キーの和をもつ数字の並びの個数とその総和を記録しておく。

import Data.List (unfoldr)
import Data.Map (Map, map, empty, singleton, mapKeysMonotonic,
                 unionsWith, foldWithKey, findWithDefault)
import Prelude hiding(map)

type State = Map Int (Integer, Integer)

append :: Int -> Int -> State -> State
append k d = mapKeysMonotonic (+d).map add
    where add (s, t) = (s, t + toInteger d*s*10^k)

next :: Int -> State -> State
next k m = unionsWith add [append k d m | d <- [0..9]]
    where add (s1, t1) (s2, t2) = (s1+s2, t1+t2)

count :: Int -> State -> State -> Integer
count k m = foldWithKey count' 0
    where count' v (s, t) c
              | even k    = c + s*(t-t0)*10^k2 + (s-s0)*t
              | otherwise = c + s*(t-t0)*10^(k2+2) + s*(s-s0)*45*10^k2 + (s-s0)*t*10
              where (s0, t0) = findWithDefault (0, 0) v m
                    k2 = div k 2

p217 :: Int -> Integer
p217 n = (`mod` (3^15)).sum.take n.unfoldr step $ (1, empty, singleton 0 (1, 0))
    where step (k, m0, m1) = Just (count k m0' m1', (k+1, m0', m1'))
              where m0' = if even k then m1 else m0
                    m1' = if even k then  next (div k 2 - 1) m1 else m1

場合分けが多いのが気にくわない。
whereも多いし。