nim-lib/nt.nim

import std/[sequtils]

const MOD* = 1000000007
type ll* = int64

proc reduce*(x: ll): ll =
  var r = x mod MOD
  if r < 0: r += MOD
  r

proc bexp*(b: ll, p: ll): ll =
  var
    curr = reduce(b)
    power = p
    res: ll = 1
  while power > 0:
    if (power and 1) == 1:
      res = (res * curr) mod MOD
    curr = (curr * curr) mod MOD
    power = power shr 1
  res

proc inv*(x: ll): ll =
  bexp(x, MOD - 2)

type
  Combinatorics*[N: static[int]] = object
    factorialCache: array[N+1, ll]
    inverseFactorialCache: array[N+1, ll]

proc initCombinatorics*[N: static[int]](): Combinatorics[N] =
  var c: Combinatorics[N]
  c.factorialCache[0] = 1

  for i in 1..N:
    c.factorialCache[i] = (c.factorialCache[i-1] * i) mod MOD

  c.inverseFactorialCache[N] = inv(c.factorialCache[N])

  for i in countdown(N-1, 0):
    c.inverseFactorialCache[i] =
      ((i+1).ll * c.inverseFactorialCache[i+1]) mod MOD

  c

proc fact*[N](c: Combinatorics[N], x: int): ll =
  c.factorialCache[x]

proc ifact*[N](c: Combinatorics[N], x: int): ll =
  c.inverseFactorialCache[x]

proc permute*[N](c: Combinatorics[N], n, k: int): ll =
  if k > n: return 0
  (c.fact(n) * c.ifact(n-k)) mod MOD

proc choose*[N](c: Combinatorics[N], n, k: int): ll =
  if k > n: return 0
  (((c.fact(n) * c.ifact(k)) mod MOD) * c.ifact(n-k)) mod MOD

proc gcd*(a, b: ll): ll =
  if b == 0: a else: gcd(b, a mod b)

proc lcm*(a, b: ll): ll =
  (a div gcd(a,b)) * b

type
  EGCDResult* = object
    x*: ll
    y*: ll
    d*: ll

proc extendedGcd*(a, b: ll): EGCDResult =
  if b == 0:
    return EGCDResult(x: 1, y: 0, d: a)

  let next = extendedGcd(b, a mod b)

  EGCDResult(
    x: next.y,
    y: next.x - (a div b) * next.y,
    d: next.d
  )

type
  FactorSieve*[N: static[int]] = object
    sieve: array[N+1, int]

proc initFactorSieve*[N: static[int]](): FactorSieve[N] =
  var s: FactorSieve[N]
  s.sieve[1] = 1

  for i in 2..N:
    if s.sieve[i] == 0:
      for j in countup(i, N, i):
        s.sieve[j] = i

  s

proc factor*[N](s: FactorSieve[N], x0: int): seq[(int,int)] =
  var x = x0
  var res: seq[(int,int)] = @[]

  while x != 1:
    let curr = s.sieve[x]
    var cnt = 0

    while s.sieve[x] == curr:
      inc cnt
      x = x div curr

    res.add((curr,cnt))

  res

type
  BooleanSieve*[N: static[int]] = object
    isComposite: array[N+1, bool]

proc initBooleanSieve*[N: static[int]](): BooleanSieve[N] =
  var s: BooleanSieve[N]
  s.isComposite[1] = true

  for i in 2..N:
    if not s.isComposite[i]:
      for j in countup(2*i, N, i):
        s.isComposite[j] = true

  s

proc isPrime*[N](s: BooleanSieve[N], x: int): bool =
  not s.isComposite[x]

# ---------------- MAIN ----------------

proc main() =
  let comb = initCombinatorics[100000]()
  echo "10 choose 3 = ", comb.choose(10,3)

  echo "gcd(48,18) = ", gcd(48,18)
  echo "lcm(48,18) = ", lcm(48,18)

  let eg = extendedGcd(48,18)
  echo "extended gcd: x=", eg.x, " y=", eg.y, " d=", eg.d

  let sieve = initFactorSieve[100000]()
  echo "Factors of 84: ", sieve.factor(84)

  let bs = initBooleanSieve[100000]()
  echo "Is 97 prime? ", bs.isPrime(97)

when isMainModule:
  main()