push

nt.nim

@@ -0,0 +1,152 @@
+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()