CLI helper for discrete math formulas

discrete.py

#!/usr/bin/env python

from functools import reduce
from math import e, factorial
from operator import mul
from os.path import basename
from sys import argv

# https://stackoverflow.com/a/4941932
def ncr(n, r):
    '''Calculate C(n, r) = n! / (r! * (n-r)!)'''
    if r > n:
        raise ValueError("cannot choose r=%d items from n=%d choices" % (r, n))
    r = min(r, n-r)
    numerator = reduce(mul, range(n, n-r, -1), 1)
    denominator = factorial(r)
    # when the numbers get large enough, we need arbitrary precision
    # integers to avoid overflow
    return numerator // denominator

def sterling(n, k):
    '''Calculate a sterling number of the second kind:
       the number of ways to partition an n-element set into k non-empty sets'''
    if n < 1 or k > n:
        return 0
    if n == k or k == 1:
        return 1
    return sterling(n - 1, k - 1) + k*sterling(n - 1, k)

def partition(n, k):
    '''Calculate the ways to partition a number n into k positive terms'''
    if k <= 0 or k > n:
        return 0
    if k == n:
        return 1
    return partition(n-k, k) + partition(n-1, k-1)

def gcd(n, d):
    '''Calculate the greatest common divisor of n and d using the Euclidean algorithm'''
    if d == 0:
        return n
    return gcd(d, n % d)

def _lcm(n, k):
    '''Find the least common multiple of n and k'''
    return n*k // gcd(n, k)

def lcm(n, *others):
    '''Find the least common multiple of n and arbitrarily many other numbers'''
    for k in others:
        n = _lcm(n, k)
    return n

def order(k, n):
    '''Given an element k in the group Z mod n, find the order of k'''
    return lcm(k, n) // k

def cyclic(*orders):
    '''Given the order of an arbitrary number of groups G_i,
       find whether their direct product G is cyclic'''
    # Theorem: G is cyclic <=> order(G_i) and order(G_j) are relatively prime
    # for all i != j
    # Theorem: a and b are relatively prime <=> gcd(a, b) == 1
    for (i, order_i) in enumerate(orders):
        for (j, order_j) in enumerate(orders):
            if i != j and gcd(order_i, order_j) != 1:
                return False
    return True

def derangements(n):
    '''Given a set S of n elements, calculate the number of permutations of S
    such that no element remains in its original position'''
    # NOTE: this is slower but more accurate than n! / e
    acc = 0
    i = 0
    while i <= n:
        acc += (-1)**i / factorial(i)
        i += 1
    return int(acc * factorial(n))

def print_generated_group(i, n):
    '''Given an element i in Z mod n, return the subgroup <i> generated by i'''
    group = set()
    j = i
    while j not in group:
         group.add(j)
         j = (j+i)%n
    return group

def elements_with_order(m, n):
    '''Return all elements in Z mod n which have order m'''
    return [i for i in range(1, n) if n/gcd(i, n) == m]

cli_opts = {
        'combinations': (ncr, 2),
        'ncr': (ncr, 2),
        'C': (ncr, 2),
        'c': (ncr, 2),

        'sterling': (sterling, 2),
        'S': (sterling, 2),
        's': (sterling, 2),

        'factorial': (factorial, 1),
        'f': (factorial, 1),

        'partition': (partition, 2),
        'P': (partition, 2),
        'p': (partition, 2),

        'euclidean': (gcd, 2),
        'gcd': (gcd, 2),
        'g': (gcd, 2),

        'lcm': (lcm, None),

        'order': (order, 2),
        'o': (order, 2),

        'cyclic': (cyclic, None),

        'derangements': (derangements, 1),
        'D': (derangements, 1),
        'd': (derangements, 1),
}

var_names = ['n', 'k', 'd', 'm', 'a', 'b', 'c',]

if __name__ == '__main__':
    (func, arity) = cli_opts[basename(argv[0])]
    if arity and len(argv) != arity + 1:
        exit("usage: " + ' '.join(argv[0:1] + var_names[:arity]))
    args = map(int, argv[1:])
    print(func(*args))

to use, put it in your $PATH as ncr, partition, or sterling