{-# LANGUAGE BangPatterns #-} {-# OPTIONS_GHC -O2 #-} {- It has been more than seven years since we solved a beautiful problem, enumerating lucky numbers. This series corresponds to the entry A000959 in the numerical series database OEIS. I hadn't realized that, at least among those listed there, ours is by far the fastest implementation for generating them. So we will revisit the algorithm and suggest an entry in **OEIS**. ## Lucky Numbers and Accumulated Cycles Generating lucky numbers consists of, starting from a certain number, making skips while discarding numbers. You start with all natural numbers, from position 1 with skips of 2. ```text V 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 ... * * * * * * * * ``` This is the same as saying we start at 1 and skip according to the cyclic set `{2}`. After jumping from 1 using those cyclic skips, we get: ```text V 1 3 5 7 9 11 13 15 17 ... ``` And we will skip every 3, meaning *the last time we skipped 2 we actually skipped 2+3*, resulting in: ```text V 1 3 5 7 9 11 13 15 17 ... * * * ``` But if we look closely, it means that for every three groups of 2, we skip another number. That is, the cycle is `{2, 4}`. Why? The previous cycle was `{2}`, which has a size of 1. Since every 3 there is another number to eliminate, we must repeat the cycle 3 times, resulting in `{2, 2, 2}`. However, that third number must be eliminated, leaving the cycle as `{2, 4}`. In other words, we start with `i = 1` (the index of the lucky number, with `i = 1, 2, 3, 4, ...`), initially `n{1} = 1` and `C{2} = {2}`, so the next lucky number is `n{i} = n{i-1} + C{i-1}{i-1}`, in this case `n{2} = n{1} + C{1}{1} = 1 + 2 = 3`. To calculate `C{i}`, we need to expand `n{i}` times `C{i-1}` and merge the skips at every `n{i}`. For `i = 2`, `C{2} = {2, 4}` as we have seen. Then, for `i = 3`, `n{i} = n{i-1} + C{i-1}{i-1} = n{2} + C{2}{2} = 3 + 4 = 7`. Thus: ```text C{3} = {2, 4, 2, 4, 2, 4 +2, 4, 2, 4, 2, 4, 2 +4} 1 2 3 4 5 6 7 8 9 10 11 12 13 14 ^ ^ C{3} = {2, 4, 2, 4, 2, 6, 4, 2, 4, 2, 4, 6} ``` For `i = 4`, `n{4} = 7 + 2 = 9`, etc. ## Generating the Next Lucky Number Given a cycle `C = {2, 4, 2, ... }`, the `i-th` lucky number consists of adding 1 from the left `i` times. ## Updating the Cycle of Increments Given a cycle `C` and the next lucky number `n`, we must repeat `C` `n` times and merge the increments at multiples of `n` with their predecessors (as we did in the earlier examples). However, it can be optimized. For example, `C{3}` has a size of 12, and `n{4} = 9`, so in theory the next size `C{4}` should be `12 * 9 = 108`. But since `LCM(12, 9)` is 36, this means the same cycle repeats 3 times. Therefore, instead of expanding to 108, it suffices to do so up to 36. ## Optimization in Haskell or Okasaki Structures In addition to optimizing the cycle length, in Haskell or any Okasaki structure, the expansions of an increment list cloned multiple times will not take up additional memory. ## Finger Tree When we have a cycle `C`, we must efficiently reach the `i-th` position. If the skips in `C` are the "fingers" and the intermediate nodes are the accumulated values, representing `C` as a finger tree seems appropriate. ## Zipper When we need to update `C` by traversing it to merge increments at certain positions, we must act on "nearby" positions without having to rebuild the entire tree from the fingers (positions we merge) to the root. An efficient way to do this is to use a zipper tree. ## Limiting Expansion It is not always necessary to fully expand cycle `C`. With it, we can generate numbers, provided the increment should not be merged. For `C{i-1}`, this occurs at position `i-2` (i.e., increments in `C{i-1}` above position `n{i}-2` might need adjustment). For example, `C{3}` has a size of 12, but `n{4}=9`, so we can only generate, with `C{3}`, another `n{4}-2 = 9-2 = 7` lucky numbers. Conversely, if we want to generate lucky numbers up to `K`, we must expand until `n{i} - 2 >= K`. ## Efficiency Comparison O(n^2) David W. Wilson https://oeis.org/A000959/a000959.txt O(n log n) Using Zipped Finger Tree ```text n-th lucky-n O(n^2) O(n log n) 71918 999987 0.821 1.281 136412 1999983 2.678 3.204 198665 2999967 5.420 5.553 259469 3999999 9.035 7.925 319305 4999995 13.647 9.980 378356 5999979 19.086 13.268 436814 6999999 25.241 15.520 494718 8000001 32.152 17.909 552174 8999997 39.748 20.686 609237 9999997 48.110 23.351 ``` Tested up to the 5286238-th lucky number 99999979. $ time ./lucky 100000000 > lucky.100m.haskell real 6m28,466s $ ( time ./david.wilson 5286238; echo "") > lucky.100m.ccode real 51m11,174s $ diff lucky.100m.ccode lucky.100m.haskell -} module Main (main) where import System.Environment (getArgs) -- Tree data type data F = N { fN_ :: !Int -- number of elements in this subtree , fL :: !F , fR :: !F } | L { fV :: !Int } -- leaf deriving Show -- Zipper holes data H = LH !Int !F | RH !Int !F -- Zipper structure data Z = Z !Int !F [H] -- Merges two subtrees {-# INLINE (.^.) #-} (.^.) :: F -> F -> F l .^. r = N (fN l + fN r) l r -- Returns the size of the tree {-# INLINE fN #-} fN :: F -> Int fN (N n _ _) = n fN (L _) = 1 -- Indexing: 1-based access {-# INLINE index #-} index :: F -> Int -> Int index (L n) 1 = n index (N _ l r) i = let !ln = fN l in if i <= ln then index l i else index r (i - ln) -- Initialize zipper from a tree {-# INLINE zipper #-} zipper :: F -> Z zipper m = Z 0 m [] -- Extract the current tree from the zipper {-# INLINE unzipper #-} unzipper :: Z -> F unzipper (Z _ m _) = m -- Move right in the zipper {-# INLINE goRight #-} goRight :: Z -> Z goRight (Z a (N _ l r) zs) = Z (a + fN l) r (RH a l : zs) goRight z = z -- Move left in the zipper {-# INLINE goLeft #-} goLeft :: Z -> Z goLeft (Z a (N _ l r) zs) = Z a l (LH a r : zs) goLeft z = z -- Move up in the zipper {-# INLINE goUp #-} goUp :: Z -> Z goUp (Z _ r (RH a l : zs)) = Z a (l .^. r) zs goUp (Z _ l (LH a r : zs)) = Z a (l .^. r) zs goUp z@(Z _ _ []) = z -- Traverse the tree via zipper until a condition is met, then apply f {-# INLINE goto #-} goto :: (Int -> Z -> Z) -> Int -> Z -> Z goto f i = go where go z@(Z a (N s l r) zs) | a < i && i <= a + fN l = go (goLeft z) | a < i && i <= a + s = go (goRight z) | null zs = z | otherwise = go (goUp z) go z = f i z -- Adjust the tree by merging at intervals of d {-# INLINE adjust #-} adjust :: Int -> F -> F adjust d m = let !lasti = let nm = fN m in nm - (nm `mod` d) doRemove i (Z _ (L n) (RH a l : zs)) = goto (doAdd n) (i - 1) $ goUp (Z a l zs) doRemove i (Z _ (L n) (LH a r : zs)) = goto (doAdd n) (i - 1) $ goUp (Z (a - 1) r zs) doRemove i z = goto doRemove i z doAdd n' i (Z a (L n) zs) = goto doRemove (i - d + 1) $ goUp $ Z a (L (n + n')) zs in unzipper (doRemove lasti (zipper m)) -- Repeat a tree n times using a "power of 2" strategy {-# INLINE repeatF #-} repeatF :: Int -> F -> F repeatF n f | n <= 0 = error "repeatF called with non-positive multiplier" | n == 1 = f | otherwise = let (d, r) = n `divMod` 2 t = repeatF d f in case r of 0 -> t .^. t _ -> t .^. t .^. f -- Expand tree without exceeding k elements {-# INLINE expand #-} expand :: Int -> Int -> F -> F expand k n f = let s = fN f r = lcm s n `div` s maxRepeat = min (k `div` s) r in adjust n (repeatF maxRepeat f) -- Generate lucky numbers {-# INLINE lucky #-} lucky :: Int -> [Int] lucky k = let -- r: building expansions r !n !f !i | n - 2 >= k = a n f i | otherwise = let !n' = n + index f i in n : r n' (expand k n' f) (i + 1) -- a: once expansions end, just collect a !n !f !i | n > k = [] | otherwise = let !n' = n + index f i in n : a n' f (i + 1) in r 1 (L 2) 1 main :: IO () main = do [k] <- map read <$> getArgs let result = lucky k mapM_ (\(i, val) -> putStrLn (show i ++ " " ++ show val)) (zip [1..] result) putStrLn "" {- $ stack ghc -- -O2 l2.hs [1 of 2] Compiling Main ( l2.hs, l2.o ) [2 of 2] Linking l2 # generate up to 51 lucky number $ ./l2 51 1 1 2 3 3 7 4 9 5 13 6 15 7 21 8 25 9 31 10 33 11 37 12 43 13 49 14 51 -}