{-# LANGUAGE FlexibleContexts #-}

-- |
-- Module      : Math.OEIS
-- Copyright   : (c) Eric Bailey, 2020-2024
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- Portability : POSIX
--
-- Some sequences from [OEIS](https://oeis.org)
module Math.OEIS
  ( a000040,
    a000078,
    a000111,
    a000182,
    a000217,
    a000290,
    a000364,
    a002378,
    a003313,
    a007947,
    a007953,
    a027748,
    a034705,
    a051885,
    a056924,
    a060735,
    a076314,
    a082949,
    a111251,
    a204692,
    a211264,
    digitSum,
    distinctPrimeFactors,
    squares,
    triangularNumbers,
  )
where

import Control.Arrow (first, (***), (>>>))
import Control.Monad (ap, join)
import Data.Bits (Bits)
import Data.Foldable (toList)
import qualified Data.IntSet as IS
import Data.List.Infinite (Infinite ((:<)), (...), (....))
import qualified Data.List.Infinite as Infinite
import Data.List.NonEmpty (NonEmpty ((:|)))
import qualified Data.List.NonEmpty as NE
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Ratio (numerator)
import Data.Set (deleteFindMin, insert, singleton)
import Math.Combinatorics.Exact.Binomial (choose)
import Math.NumberTheory.ArithmeticFunctions (tau)
import Math.NumberTheory.Primes (Prime, UniqueFactorisation, factorise, nextPrime, unPrime)
import Math.NumberTheory.Primes.Testing (isPrime)
import Math.NumberTheory.Recurrences (bernoulli, euler)
import Math.OEIS.HybridInteger (HybridInteger (..))
import Numeric.Natural (Natural)

-- | The prime numbers, i.e. [A000040](https://oeis.org/A000040).
a000040 :: (Bits a, Enum (Prime a), Integral a, UniqueFactorisation a) => Infinite a
a000040 :: forall a.
(Bits a, Enum (Prime a), Integral a, UniqueFactorisation a) =>
Infinite a
a000040 = (Prime a -> a) -> Infinite (Prime a) -> Infinite a
forall a b. (a -> b) -> Infinite a -> Infinite b
Infinite.map Prime a -> a
forall a. Prime a -> a
unPrime (a -> Prime a
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
nextPrime a
2 ...)

-- | Tetranacci numbers, i.e. [A000078](https://oeis.org/A000078).
a000078 :: (Integral a) => Infinite a
a000078 :: forall a. Integral a => Infinite a
a000078 = ((a, a, a, a) -> (a, (a, a, a, a))) -> (a, a, a, a) -> Infinite a
forall b a. (b -> (a, b)) -> b -> Infinite a
Infinite.unfoldr (a, a, a, a) -> (a, (a, a, a, a))
forall {d}. Num d => (d, d, d, d) -> (d, (d, d, d, d))
go (a
0, a
0, a
0, a
1)
  where
    go :: (d, d, d, d) -> (d, (d, d, d, d))
go (d
a, d
b, d
c, d
d) = (d
a, (d
b, d
c, d
d, d
a d -> d -> d
forall a. Num a => a -> a -> a
+ d
b d -> d -> d
forall a. Num a => a -> a -> a
+ d
c d -> d -> d
forall a. Num a => a -> a -> a
+ d
d))

-- | Euler up/down numbers: e.g.f. \(\sec(x) + \tan(x)\), i.e.
-- [A000111](https://oeis.org/A000111).
a000111 :: (Integral a) => Infinite a
a000111 :: forall a. Integral a => Infinite a
a000111 = Infinite a -> Infinite a -> Infinite a
forall a. Infinite a -> Infinite a -> Infinite a
Infinite.interleave Infinite a
forall a. Integral a => Infinite a
a000364 Infinite a
forall a. Integral a => Infinite a
a000182

-- | Tangent (or \"Zag\") numbers: e.g.f. \(\tan(x)\), also (up to signs) e.g.f.
-- \(\tanh(x)\), i.e. [A000182](https://oeis.org/A000182).
a000182 :: (Integral a) => Infinite a
a000182 :: forall a. Integral a => Infinite a
a000182 =
  ((a, Ratio a) -> a) -> Infinite (a, Ratio a) -> Infinite a
forall a b. (a -> b) -> Infinite a -> Infinite b
Infinite.map
    (\(a
n, Ratio a
b) -> Ratio a -> a
forall a. Ratio a -> a
numerator (Ratio a
2 Ratio a -> a -> Ratio a
forall a b. (Num a, Integral b) => a -> b -> a
^ a
n Ratio a -> Ratio a -> Ratio a
forall a. Num a => a -> a -> a
* (Ratio a
2 Ratio a -> a -> Ratio a
forall a b. (Num a, Integral b) => a -> b -> a
^ a
n Ratio a -> Ratio a -> Ratio a
forall a. Num a => a -> a -> a
- Ratio a
1) Ratio a -> Ratio a -> Ratio a
forall a. Num a => a -> a -> a
* Ratio a -> Ratio a
forall a. Num a => a -> a
abs Ratio a
b) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
n)
    (Infinite (a, Ratio a) -> Infinite a)
-> Infinite (a, Ratio a) -> Infinite a
forall a b. (a -> b) -> a -> b
$ Infinite a -> Infinite (Ratio a) -> Infinite (a, Ratio a)
forall a b. Infinite a -> Infinite b -> Infinite (a, b)
Infinite.zip ((a
2, a
4) ....)
    (Infinite (Ratio a) -> Infinite (a, Ratio a))
-> Infinite (Ratio a) -> Infinite (a, Ratio a)
forall a b. (a -> b) -> a -> b
$ Infinite (Ratio a) -> Infinite (Ratio a)
forall a. Infinite a -> Infinite a
Infinite.tail
    (Infinite (Ratio a) -> Infinite (Ratio a))
-> Infinite (Ratio a) -> Infinite (Ratio a)
forall a b. (a -> b) -> a -> b
$ Infinite (Ratio a) -> Infinite (Ratio a)
forall a. Infinite a -> Infinite a
everyOther Infinite (Ratio a)
forall a. Integral a => Infinite (Ratio a)
bernoulli

-- | Triangular numbers, i.e. [A000217](https://oeis.org/A000217).
a000217 :: (Integral a) => a -> a
a000217 :: forall a. Integral a => a -> a
a000217 a
n = a
n a -> a -> a
forall a. Num a => a -> a -> a
* (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2

-- | See 'a000217'.
triangularNumbers :: (Integral a) => Infinite a
triangularNumbers :: forall a. Integral a => Infinite a
triangularNumbers = (a -> a) -> Infinite a -> Infinite a
forall a b. (a -> b) -> Infinite a -> Infinite b
Infinite.map a -> a
forall a. Integral a => a -> a
a000217 (a
0 ...)

a000290, squares :: (Enum a, Num a) => Infinite a
a000290 :: forall a. (Enum a, Num a) => Infinite a
a000290 = (a -> a -> a) -> a -> Infinite a -> Infinite a
forall b a. (b -> a -> b) -> b -> Infinite a -> Infinite b
Infinite.scanl a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 ((a
1, a
3) ....)
-- ^ The squares: \(a000290(n) = n^2\), i.e. [A000290](https://oeis.org/A000290).
squares :: forall a. (Enum a, Num a) => Infinite a
squares = Infinite a
forall a. (Enum a, Num a) => Infinite a
a000290
-- ^ See 'a000290'.

-- | Euler (or secant or \"Zig\") numbers: e.g.f. (even powers only)
-- \(\sec(x) = \frac{1}{\cos(x)}\), i.e. [A000364](https://oeis.org/A000364).
a000364 :: (Integral a) => Infinite a
a000364 :: forall a. Integral a => Infinite a
a000364 = Infinite a -> Infinite a
forall a. Infinite a -> Infinite a
everyOther ((a -> a) -> Infinite a -> Infinite a
forall a b. (a -> b) -> Infinite a -> Infinite b
Infinite.map a -> a
forall a. Num a => a -> a
abs Infinite a
forall a. Integral a => Infinite a
euler)

-- | Oblong (or promic, pronic, or heteromecic) numbers:
-- \(a002378(n) = n(n+1)\), i.e. [A002378](https://oeis.org/A002378).
a002378 :: (Integral a) => a -> a
a002378 :: forall a. Integral a => a -> a
a002378 a
n = a
n a -> a -> a
forall a. Num a => a -> a -> a
* (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)

-- | Length of shortest addition chain for \(n\), i.e.
-- [A003313](https://oeis.org/A003313).
a003313 :: Natural -> Int
a003313 :: Natural -> Int
a003313 Natural
n
  | Natural
n Natural -> [Natural] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Natural
77, Natural
154] = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  | Bool
otherwise = Int
k
  where
    k :: Int
k = [Natural] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Natural -> [Natural]
powerTree Natural
n) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | Largest squarefree number diving \(n\): the square free kernel of \(n\),
-- \(\text{rad}(n)\), radical of \(n\), i.e. [A007947](https://oeis.org/A007947).
a007947 :: (UniqueFactorisation a) => a -> a
a007947 :: forall a. UniqueFactorisation a => a -> a
a007947 = NonEmpty a -> a
forall a. Num a => NonEmpty a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (NonEmpty a -> a) -> (a -> NonEmpty a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> NonEmpty a
forall a. UniqueFactorisation a => a -> NonEmpty a
distinctPrimeFactors

-- | The [distinct prime factors](https://mathworld.wolfram.com/DistinctPrimeFactors.html)
-- of a given number.
--
-- >>> distinctPrimeFactors 504
-- 2 :| [3,7]
--
-- See also 'a007947' and 'a027748'.
distinctPrimeFactors :: (UniqueFactorisation a) => a -> NonEmpty a
distinctPrimeFactors :: forall a. UniqueFactorisation a => a -> NonEmpty a
distinctPrimeFactors = ((Prime a, Word) -> a) -> NonEmpty (Prime a, Word) -> NonEmpty a
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NE.map (Prime a -> a
forall a. Prime a -> a
unPrime (Prime a -> a)
-> ((Prime a, Word) -> Prime a) -> (Prime a, Word) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Prime a, Word) -> Prime a
forall a b. (a, b) -> a
fst) (NonEmpty (Prime a, Word) -> NonEmpty a)
-> (a -> NonEmpty (Prime a, Word)) -> a -> NonEmpty a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Prime a, Word)] -> NonEmpty (Prime a, Word)
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList ([(Prime a, Word)] -> NonEmpty (Prime a, Word))
-> (a -> [(Prime a, Word)]) -> a -> NonEmpty (Prime a, Word)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [(Prime a, Word)]
forall a. UniqueFactorisation a => a -> [(Prime a, Word)]
factorise

a007953, digitSum :: (Integral a) => a -> a
a007953 :: forall a. Integral a => a -> a
a007953 a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
10 = a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
100 = a -> a
forall a. Integral a => a -> a
a076314 a
n
  | Bool
otherwise = (a -> a) -> (a, a) -> (a, a)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first a -> a
forall a. Integral a => a -> a
a007953 ((a, a) -> (a, a)) -> ((a, a) -> a) -> (a, a) -> a
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(+) ((a, a) -> a) -> (a, a) -> a
forall a b. (a -> b) -> a -> b
$ a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
divMod a
n a
10
-- ^ Digital sum of \(n\), i.e. [A007953](https://oeis.org/A007953).
digitSum :: forall a. Integral a => a -> a
digitSum = a -> a
forall a. Integral a => a -> a
a007953
-- ^ See 'a007953'.

-- | Irregular triangle in which first row is \(1\), \(n\)-th row \((n > 1)\)
-- lists distinct prime factors of \(n\), i.e.
-- [A027748](https://oeis.org/A027748).
--
-- See 'distinctPrimeFactors'.
a027748 :: (Enum a, UniqueFactorisation a) => Infinite a
a027748 :: forall a. (Enum a, UniqueFactorisation a) => Infinite a
a027748 = a
1 a -> Infinite a -> Infinite a
forall a. a -> Infinite a -> Infinite a
:< (a -> NonEmpty a) -> Infinite a -> Infinite a
forall a b. (a -> NonEmpty b) -> Infinite a -> Infinite b
Infinite.concatMap a -> NonEmpty a
forall a. UniqueFactorisation a => a -> NonEmpty a
distinctPrimeFactors (a
2 ...)

-- | Numbers that are sums of consecutive squares, i.e.
-- [A034705](https://oeis.org/A034705).
--
-- See also 'a000290' and 'squares'.
a034705 :: Infinite Int
a034705 :: Infinite Int
a034705 = ((Int, Infinite (NonEmpty Int), IntSet)
 -> (Int, (Int, Infinite (NonEmpty Int), IntSet)))
-> (Int, Infinite (NonEmpty Int), IntSet) -> Infinite Int
forall b a. (b -> (a, b)) -> b -> Infinite a
Infinite.unfoldr (Int, Infinite (NonEmpty Int), IntSet)
-> (Int, (Int, Infinite (NonEmpty Int), IntSet))
go (Int
0, Infinite (NonEmpty Int) -> Infinite (NonEmpty Int)
forall a. Infinite a -> Infinite a
Infinite.tail (Infinite Int -> Infinite (NonEmpty Int)
forall a. Infinite a -> Infinite (NonEmpty a)
Infinite.inits1 Infinite Int
forall a. (Enum a, Num a) => Infinite a
a000290), [Int] -> IntSet
IS.fromList [Int
0])
  where
    go :: (Int, Infinite (NonEmpty Int), IntSet)
-> (Int, (Int, Infinite (NonEmpty Int), IntSet))
go (Int
x, NonEmpty Int
vs :< Infinite (NonEmpty Int)
vss, IntSet
seen)
      | Int
minSeen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x = (Int
minSeen, (Int
x, NonEmpty Int
vs NonEmpty Int -> Infinite (NonEmpty Int) -> Infinite (NonEmpty Int)
forall a. a -> Infinite a -> Infinite a
:< Infinite (NonEmpty Int)
vss, IntSet
seen'))
      | Bool
otherwise = (Int, Infinite (NonEmpty Int), IntSet)
-> (Int, (Int, Infinite (NonEmpty Int), IntSet))
go (Int
w, Infinite (NonEmpty Int)
vss, IntSet -> IntSet -> IntSet
IS.union IntSet
seen ([Int] -> IntSet
IS.fromList ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
w [Int]
ws)))
      where
        (Int
w :| [Int]
ws) = NonEmpty Int -> NonEmpty Int
forall a. NonEmpty a -> NonEmpty a
NE.reverse NonEmpty Int
vs
        (Int
minSeen, IntSet
seen') = IntSet -> (Int, IntSet)
IS.deleteFindMin IntSet
seen

-- | Smallest number whose sum of digits is \(n\), i.e.
-- [A051885](https://oeis.org/A051885).
--
-- See also 'a007953' and 'digitSum'.
a051885 :: (Integral a) => a -> a
a051885 :: forall a. Integral a => a -> a
a051885 = (a -> a -> (a, a)) -> a -> a -> (a, a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
divMod a
9 (a -> (a, a)) -> ((a, a) -> a) -> a -> a
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> ((a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (a -> a) -> (a -> a) -> (a, a) -> (a, a)
forall b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** (a
10 ^) ((a, a) -> (a, a)) -> ((a, a) -> a) -> (a, a) -> a
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(*) ((a, a) -> a) -> (a -> a) -> (a, a) -> a
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> a -> a -> a
forall a. Num a => a -> a -> a
subtract a
1)

-- | Number of divisors of \(n\) that are smaller than \(\sqrt{n}\), i.e.
-- [A056924](https://oeis.org/A056924).
--
-- \(a056924(n) = N(4n)\) from [Project Euler Problem 174: Hollow Square
-- Laminae II](https://projecteuler.net/problem=174).
a056924 :: (Integral a, UniqueFactorisation a) => a -> a
a056924 :: forall a. (Integral a, UniqueFactorisation a) => a -> a
a056924 = (a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
forall n a. (UniqueFactorisation n, Num a) => n -> a
tau

-- | \(a060735(1) = 1\), \(a060735(2) = 2\); thereafter, \(a060735(n)\) is the
-- smallest number \(m\) not yet in the sequence such that every prime that
-- divides \(a(n-1)\) also divides \(m\), i.e.
-- [A060735](https://oeis.org/A060735).
--
-- See 'a007947'.
a060735 :: (UniqueFactorisation a) => Infinite a
a060735 :: forall a. UniqueFactorisation a => Infinite a
a060735 = (a -> a) -> a -> Infinite a
forall a. (a -> a) -> a -> Infinite a
Infinite.iterate (a -> a -> a
forall a. Num a => a -> a -> a
(+) (a -> a -> a) -> (a -> a) -> a -> a
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
`ap` a -> a
forall a. UniqueFactorisation a => a -> a
a007947) a
2

-- | \(a076314(n) = \lfloor \frac{n}{10} \rfloor + (n \bmod 10)\), i.e.
-- [A0076314](https://oeis.org/A0076314).
a076314 :: (Integral a) => a -> a
a076314 :: forall a. Integral a => a -> a
a076314 = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(+) ((a, a) -> a) -> (a -> (a, a)) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> (a, a)) -> a -> a -> (a, a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
divMod a
10

-- | Numbers of the form \(p^{q}q^{p}\), with distinct primes \(p\) and \(q\),
-- i.e. [A082949](https://oeis.org/A082949).
a082949 :: Infinite HybridInteger
a082949 :: Infinite HybridInteger
a082949 = (Set HybridInteger -> (HybridInteger, Set HybridInteger))
-> Set HybridInteger -> Infinite HybridInteger
forall b a. (b -> (a, b)) -> b -> Infinite a
Infinite.unfoldr Set HybridInteger -> (HybridInteger, Set HybridInteger)
go (HybridInteger -> Set HybridInteger
forall a. a -> Set a
singleton (Prime Integer -> Prime Integer -> HybridInteger
HybridInteger (Integer -> Prime Integer
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
nextPrime Integer
2) (Integer -> Prime Integer
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
nextPrime Integer
3)))
  where
    go :: Set HybridInteger -> (HybridInteger, Set HybridInteger)
go Set HybridInteger
seen
      | Prime Integer
p' Prime Integer -> Prime Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Prime Integer
q = (HybridInteger
hi, HybridInteger -> Set HybridInteger -> Set HybridInteger
forall a. Ord a => a -> Set a -> Set a
insert (Prime Integer -> Prime Integer -> HybridInteger
HybridInteger Prime Integer
p' Prime Integer
q) Set HybridInteger
seen')
      | Bool
otherwise = (HybridInteger
hi, Set HybridInteger
seen')
      where
        seen' :: Set HybridInteger
seen' = HybridInteger -> Set HybridInteger -> Set HybridInteger
forall a. Ord a => a -> Set a -> Set a
insert (Prime Integer -> Prime Integer -> HybridInteger
HybridInteger Prime Integer
p Prime Integer
q') Set HybridInteger
rest
        p' :: Prime Integer
p' = Prime Integer -> Prime Integer
forall a. Enum a => a -> a
succ Prime Integer
p
        q' :: Prime Integer
q' = Prime Integer -> Prime Integer
forall a. Enum a => a -> a
succ Prime Integer
q
        (hi :: HybridInteger
hi@(HybridInteger Prime Integer
p Prime Integer
q), Set HybridInteger
rest) = Set HybridInteger -> (HybridInteger, Set HybridInteger)
forall a. Set a -> (a, Set a)
deleteFindMin Set HybridInteger
seen

-- | Numbers \(k\) such that \(3k^2 + 3k + 1\) is prime, i.e.
-- [A111251](https://oeis.org/A111251).
a111251 :: Infinite Integer
a111251 :: Infinite Integer
a111251 =
  (Integer -> Bool) -> Infinite Integer -> Infinite Integer
forall a. (a -> Bool) -> Infinite a -> Infinite a
Infinite.filter Integer -> Bool
isPrime (Infinite Integer -> Infinite Integer)
-> Infinite Integer -> Infinite Integer
forall a b. (a -> b) -> a -> b
$
    (Integer -> Integer) -> Infinite Integer -> Infinite Integer
forall a b. (a -> b) -> Infinite a -> Infinite b
Infinite.map (\Integer
k -> Integer
3 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
3 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) (Integer
1 ...)

-- | The number of base-10 [bouncy numbers](https://projecteuler.net/problem=113)
-- below \(10^{n}\), i.e. [A204692](https://oeis.org/A204692).
a204692 :: (Integral a) => a -> a
a204692 :: forall a. Integral a => a -> a
a204692 a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
3 = a
0
  | Bool
otherwise = a
10 a -> a -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ a
n a -> a -> a
forall a. Num a => a -> a -> a
- (a -> a -> a
forall a. Integral a => a -> a -> a
choose (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
10) a
10 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a -> a
forall a. Integral a => a -> a -> a
choose (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
9) a
9 a -> a -> a
forall a. Num a => a -> a -> a
- a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
10 a -> a -> a
forall a. Num a => a -> a -> a
* a
n)

-- | Number of integer pairs \((x,y,)\) such that \(0 < x < y \leq n\) and
-- \(xy \leq n\), i.e. [A211264](https://oeis.org/A211264).
a211264 :: (Integral a) => a -> a
a211264 :: forall a. Integral a => a -> a
a211264 a
n = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a
n a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
k | a
k <- [a
1 .. a
m]] a -> a -> a
forall a. Num a => a -> a -> a
- a
m a -> a -> a
forall a. Num a => a -> a -> a
* (a
m a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2
  where
    m :: a
m = forall a b. (RealFrac a, Integral b) => a -> b
floor @Double (Double -> Double
forall a. Floating a => a -> a
sqrt (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n))

-- -------------------------------------------------------- [ Helper functions ]

everyOther :: Infinite a -> Infinite a
everyOther :: forall a. Infinite a -> Infinite a
everyOther (a
x :< a
_ :< Infinite a
xs) = a
x a -> Infinite a -> Infinite a
forall a. a -> Infinite a -> Infinite a
:< Infinite a -> Infinite a
forall a. Infinite a -> Infinite a
everyOther Infinite a
xs

powerTree :: Natural -> [Natural]
powerTree :: Natural -> [Natural]
powerTree Natural
n
  | Natural
n Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
<= Natural
0 = []
  | Bool
otherwise = NonEmpty (Map Natural [Natural]) -> [Natural]
forall {t :: * -> *} {a}.
Foldable t =>
NonEmpty (Map Natural (t a)) -> [a]
findIn NonEmpty (Map Natural [Natural])
levels
  where
    findIn :: NonEmpty (Map Natural (t a)) -> [a]
findIn (Map Natural (t a)
m :| [Map Natural (t a)]
ms) = [a] -> (t a -> [a]) -> Maybe (t a) -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (NonEmpty (Map Natural (t a)) -> [a]
findIn ([Map Natural (t a)] -> NonEmpty (Map Natural (t a))
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList [Map Natural (t a)]
ms)) t a -> [a]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Natural -> Map Natural (t a) -> Maybe (t a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Natural
n Map Natural (t a)
m)

levels :: NonEmpty (Map Natural [Natural])
levels :: NonEmpty (Map Natural [Natural])
levels = (Map Natural [Natural] -> Map Natural [Natural])
-> Map Natural [Natural] -> NonEmpty (Map Natural [Natural])
forall a. (a -> a) -> a -> NonEmpty a
NE.iterate Map Natural [Natural] -> Map Natural [Natural]
expand (Natural -> [Natural] -> Map Natural [Natural]
forall k a. k -> a -> Map k a
Map.singleton Natural
1 [Natural
1])
  where
    expand :: Map Natural [Natural] -> Map Natural [Natural]
expand = (Map Natural [Natural]
 -> Map Natural [Natural] -> Map Natural [Natural])
-> Map Natural [Natural] -> Map Natural [Natural]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ((Natural
 -> [Natural] -> Map Natural [Natural] -> Map Natural [Natural])
-> Map Natural [Natural]
-> Map Natural [Natural]
-> Map Natural [Natural]
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey Natural
-> [Natural] -> Map Natural [Natural] -> Map Natural [Natural]
forall {a}. (Ord a, Num a) => a -> [a] -> Map a [a] -> Map a [a]
addLevel)
    addLevel :: a -> [a] -> Map a [a] -> Map a [a]
addLevel a
n [a]
xs Map a [a]
tree = (a -> Map a [a] -> Map a [a]) -> Map a [a] -> [a] -> Map a [a]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a -> [a] -> a -> Map a [a] -> Map a [a]
forall {a}.
(Ord a, Num a) =>
a -> [a] -> a -> Map a [a] -> Map a [a]
addNext a
n [a]
xs) Map a [a]
tree [a]
xs
    addNext :: a -> [a] -> a -> Map a [a] -> Map a [a]
addNext a
n [a]
xs a
m = a -> [a] -> Map a [a] -> Map a [a]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
m) (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
m a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs)