Решение проблемы из Google Code Jam (2009.1AA: "Multi- базовое счастье ") Я придумал неудобное (с точки зрения кода) решение, и мне интересно, как его можно улучшить.
Краткое описание проблемы: найти наименьшее число больше 1, для которого итеративное вычисление суммы квадратов цифр достигает 1 для всех оснований из данного списка.
Или описание в псевдо-Haskell (код, который решил бы это, если бы elem
всегда мог работать для бесконечных списков):
solution =
head . (`filter` [2..]) .
all ((1 `elem`) . (`iterate` i) . sumSquareOfDigitsInBase)
И мое неудобное решение:
- Под неудобным я имею в виду такой код:
happy <- lift . lift . lift $ isHappy Set.empty base cur
- Запоминаю результаты функции isHappy. Использование монады State для запоминания результатов Map.
- Пытаясь найти первое решение, я не использовал
head
иfilter
(как это делает псевдо-haskell выше), потому что вычисление не является чистым (меняет состояние). Поэтому я повторил попытку, используя StateT со счетчиком и MaybeT, чтобы завершить вычисление, когда условие выполняется. - Уже внутри
MaybeT (StateT a (State b))
, если условие не выполняется для одной базы, нет необходимости проверять другие, поэтому у меня есть ещеMaybeT
в стеке для этого.
Код:
import Control.Monad.Maybe
import Control.Monad.State
import Data.Maybe
import qualified Data.Map as Map
import qualified Data.Set as Set
type IsHappyMemo = State (Map.Map (Integer, Integer) Bool)
isHappy :: Set.Set Integer -> Integer -> Integer -> IsHappyMemo Bool
isHappy _ _ 1 = return True
isHappy path base num = do
memo <- get
case Map.lookup (base, num) memo of
Just r -> return r
Nothing -> do
r <- calc
when (num < 1000) . modify $ Map.insert (base, num) r
return r
where
calc
| num `Set.member` path = return False
| otherwise = isHappy (Set.insert num path) base nxt
nxt =
sum . map ((^ (2::Int)) . (`mod` base)) .
takeWhile (not . (== 0)) . iterate (`div` base) $ num
solve1 :: [Integer] -> IsHappyMemo Integer
solve1 bases =
fmap snd .
(`runStateT` 2) .
runMaybeT .
forever $ do
(`when` mzero) . isJust =<<
runMaybeT (mapM_ f bases)
lift $ modify (+ 1)
where
f base = do
cur <- lift . lift $ get
happy <- lift . lift . lift $ isHappy Set.empty base cur
unless happy mzero
solve :: [String] -> String
solve =
concat .
(`evalState` Map.empty) .
mapM f .
zip [1 :: Integer ..]
where
f (idx, prob) = do
s <- solve1 . map read . words $ prob
return $ "Case #" ++ show idx ++ ": " ++ show s ++ "\n"
main :: IO ()
main =
getContents >>=
putStr . solve . tail . lines
У других участников, использующих Haskell, действительно были более хорошие решения, но решил проблему иначе. Мой вопрос касается небольших итеративных улучшений моего кода.