Problem 217

217 Balanced Numbers

探索範囲が広いので、全探索は無理。

また、個数ではなく、和が必要。

さて、どう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 <- [..9]]
where add (s1, t1) (s2, t2) = (s1+s2, t1+t2)
count :: Int -> State -> State -> Integer
count k m = foldWithKey count' 
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 (, ) v m
k2 = div k 2
p217 :: Int -> Integer
p217 n = (`mod` (3^15)).sum.take n.unfoldr step $ (1, empty, singleton  (1, ))
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も多いし。

コンテンツ