もう少し速いナルシシスト数計算

こちらの記事を見て、いじっていたら面白かったので自分も書きます。

上限

ナルシシスト数の存在する範囲に上限かあることはごく簡単に証明できる。

関数$f(x)$を「それぞれの数字の桁数乗の和」という計算と定義する。つまりナルシシスト数とは$f(x) = x$となるような数となる。
$n$桁の数の中で、$f(x)$を最大にするような数は、999…999という見ための、9が$n$回並んだ数で、その$f(999\dots 999) = n \times 9^n$ である。これを$g(n)$ とする。
さてところで、$n$桁の数の最小値は$10^{n-1}$である。これを$h(n)$ とする。
$g(n)$は、$n$が1増えるたびに $9\frac{n+1}{n}$（ほぼ9）倍になる。
$h(n)$は、$n$が1増えるたびに 10倍になる。つまり、いつか追い越される。
追い越されるとはどういうことかというと、ある桁数を越えると、$f(x)$がどうやっても$n$桁未満の数にしかならなくなる、つまり、ナルシシスト数が絶対に存在しなくなる。
計算すると、
$g(60) = 1.0782061799486587 \times 10^{59} > h(60) = 10^{59}$
$g(61) = 9.865586546530227 \times 10^{59} < h(61) = 10^{60}$
ということで、61桁以上の桁数のナルシシスト数は存在しえない。

Rosetta Codeのコード

トップページに

Rosetta Code is a programming chrestomathy site.

(chrestomathy
n. （普通初学者が外国語の勉強に用いる）名句選; 名句集, 選句集.
[株式会社研究社　新英和大辞典第6版])

なんて書いてあるけど、どうなんでしょう。

のコードを見てみる。

import Data.Bifunctor (second)

narcissiOfLength :: Int -> [Int]
narcissiOfLength nDigits = snd <$> go nDigits []
  where
    powers = ((,) <*> (^ nDigits)) <$> [0 .. 9]
    go n parents
      | 0 < n = go (pred n) (f parents)
      | otherwise = filter (isDaffodil nDigits . snd) parents
      where
        f parents
          | null parents = powers
          | otherwise =
            parents >>=
            (\(d, pwrSum) -> second (pwrSum +) <$> take (succ d) powers)

isDaffodil :: Int -> Int -> Bool
isDaffodil e n =
  (((&&) . (e ==) . length) <*> (n ==) . powerSum e) (digitList n)

powerSum :: Int -> [Int] -> Int
powerSum n = foldr ((+) . (^ n)) 0

digitList :: Int -> [Int]
digitList 0 = [0]
digitList n = go n
  where
    go 0 = []
    go x = rem x 10 : go (quot x 10)

--------------------------- TEST ---------------------------
main :: IO ()
main =
  putStrLn $
  fTable
    "Narcissistic decimal numbers of length 1-7:\n"
    show
    show
    narcissiOfLength
    [1 .. 7]

fTable :: String -> (a -> String) -> (b -> String) -> (a -> b) -> [a] -> String
fTable s xShow fxShow f xs =
  let rjust n c = drop . length <*> (replicate n c ++)
      w = maximum (length . xShow <$> xs)
  in unlines $
     s : fmap (((++) . rjust w ' ' . xShow) <*> ((" -> " ++) . fxShow . f)) xs

……。
Haskellらしく、っていうのは、こんな風に無意味に <*> とかを振り回すことじゃないと思うんですよ。
それをするのが必然な、抽象化しまくり高階計算しまくりなプログラムということでもなく、これではただの難読化でしょう。
ということで、解読がてら容赦なく直す。

import Data.Char
import System.CPUTime

narcissiOfLength :: Int -> [Integer]
narcissiOfLength nDigits = filter (isDaffodil nDigits) $
                           map snd $
                           iterate f powers !! pred nDigits
  where
--    uLimit = 10 ^ nDigits
    powers = [(i, fromIntegral i ^ nDigits) | i <- [0 .. 9]]
    f parents = [ (e, pwrSum1)
                | (d, pwrSum) <- parents
                , (e, a) <- take (succ d) powers
                , let pwrSum1 = pwrSum + a
--                , pwrSum1 <= uLimit -- 大きくなりすぎた値を切り捨てる枝刈り
                ]

isDaffodil :: Int -> Integer -> Bool
isDaffodil e n = e == length ds && n == sum (map (^ e) ds)
  where
    ds = map (fromIntegral . digitToInt) . show $ n

--------------------------- TEST ---------------------------

main :: IO ()
main = do
  putStrLn "Enter lb and ub:"
  [lb,ub] <-map read . words <$> getLine
  t1 <- getCPUTime
  putStrLn "Narcissistic decimal numbers of length " ++ show lb ++ "-" ++ show ub ++ ":\n"
  putStrLn $ fTable show show narcissiOfLength [lb .. ub]
  t2 <- getCPUTime
  print $ t2 - t1

fTable :: (a -> String) -> (b -> String) -> (a -> b) -> [a] -> String
fTable xShow fxShow f xs = unlines $
    map (\x -> rjust (xShow x) ++ " -> " ++ fxShow (f x)) xs
  where
    rjust s = replicate (w - length s) ' ' ++ s
    w = maximum $ map (length . xShow) xs

何をしているかというと、
powersが i=0～9についてiの桁数乗との対のリストで、parents の初期値。
これは、(今後使ってよい数字の最大値d, ここまでの桁数乗の和pwrSum) のリスト。
fstは一度小さくしたら大きくすることはできない。
fを桁数-1回繰り返し適用することで、0～9の数字を適当な回数使ってnDigits桁の数を作ったとき、数字の順序がどうあれ、そのような成分で作られる$x$の$f(x)$の値を全て作ることができる。
その総数は${}_{n+9} C_9$で、$n = 39$ のとき $\xcancel{1.2896910330970532 \times 10^{23}}$ となり、たぶんそんなリストはメモリに入りきらない。 $1.67710664\times 10^9$ で、メモリに収めるのは辛そう。
（訂正します。計算間違えてました。）

もう少し違うアプローチがあるんじゃないかな。

深さ優先探索

上のコードはいうなれば、parentsの内容で幅優先探索をしているようなもの。
同じ問題空間を深さ優先探索することで、そんなにメモリを消費せずに同じことができる。
さらに、それぞれの数字を何個使うかを、一度だけ選択するようにして、再帰は桁数回でなく数字の種類の10段で済ませる。

naoLen :: Int -> [Integer]
naoLen n = recur n [] 0 table
  where
-- d = 9～1について、k = 0～n について、k * d^n の表
    table :: [(Int, [(Int, Integer)])]
    table = [(d, zip [0..n] [0, x ..])
            | d <- [9, 8 .. 1]
            , let x = fromIntegral d ^ n]

    lb = 10 ^ pred n
    ub = 10 * lb

-- recur rest cnts acc tbl
-- rest : 必要な数字の残り桁数
-- cnts : 使用の確定した数字の列 [6,6,8,9,9,9] みたいな形
-- acc  : ここまでの合計 sum $ map (^ 桁数) cnts
-- tbl  : table のtails 先頭要素をいくつ使うかを今から選ぶ
-- 結果 : 発見した答えのリスト

    recur _ _ acc _ | ub <= acc = [] -- 大きすぎ n+1桁の数になってしまった。

    recur rest _ acc ((_, kvs):_) | acc + snd (kvs !! rest) < lb = [] -- 小さすぎ 残りの桁全てをこの数字で埋めても下限に届かない

    recur rest cnts acc ((d, kvs):tbl) = -- 正常系 再帰を深める
      [ ans
      | (k,v) <- take (succ rest) kvs
      , ans <- recur (rest - k) (replicate k d ++ cnts) (acc + v) tbl
      ]

    recur rest cnts acc [] = [acc | cnts == cnts1] -- tableが尽きたら、0を除いて数字の個数が一致したらナルシシスト数
      where
        cnts1 = sort $ map digitToInt $ filter ('0' /=) $ show acc

CPUTimeによる計測で、
1～20 : 17,437,500,000,000 ps = 17秒
20～39 : 4,980,390,625,000,000 ps = 83分
ということで、1時間半足らずで全てのナルシシスト数を算出できた。

環境：
Windows11, 11th Gen Core i7-1165G7 @ 2.80GHz
ghc-9.4.8

ちなみに、下限と上限で枝刈りをしているので、naoLen 61 とかはすぐに返ってくる。

おわり。

冒頭のプログラムがそのまま読めないと、若いHaskellerから「staticおじさん」ならぬ「一階おじさん」とか陰口叩かれてしまうのだろうか。