Implementations of some interesting algorithms in Python.

problems.py

# -----------------------------------------------
# Algorithms for solving string matching problems
# -----------------------------------------------

def automaton(p, abc):
    """Builds DFA using pattern p and alphabet abc."""
    m = len(p)
    d = {}
    for q in range(m+1):
        for a in abc:
            k = min(m, q+1)
            while not (p[:q] + a).endswith(p[:k]):
                k -= 1
            d[q, a] = k
    return d


def automaton_matcher(t, p, d=None):
    """Matches given pattern p in the text t. Yields all occurences."""
    if d is None:
        d = automaton(p, set(a for a in t))
    q = 0
    n, m = len(t), len(p)
    for i in range(n):
        q = d[q, t[i]]
        if q == m:
            yield i - m + 1


def prefix(p):
    r = [0] * len(p)
    k = 0
    for q in range(1, len(p)):
        while k > 0 and p[k] != p[q]:
            k = r[k-1]
        if p[k] == p[q]:
            k += 1
        r[q] = k
    return r


def kmp(t, p):
    """Implementation of the Knuth-Morris-Pratt algorithm."""
    r = prefix(p)
    q = 0
    for i in range(len(t)):
        while q > 0 and p[q] != t[i]:
            q = r[q-1]
        if p[q] == t[i]:
            q += 1
        if q == len(p):
            yield i - len(p) + 1
            q = r[q-1]
        
        
# -------------------------------------
# Algorithms for solving graph problems
# -------------------------------------

def initialize(G, s):
    """Initialize graph G and vertex s."""
    V, E = G
    d = {v: float('inf') for v in V}
    p = {v: None for v in V}
    d[s] = 0
    return d, p


def bellman_ford(G, w, s):
    """Bellman-Ford's algorithm for shortest-path search. Expects oriented graph.
    Parameter G is a graph represented as a tuple of vertexes and edges. The
    parametr w represents weights of each edge of the graph G (w: E -> R).

    Example::
    
        >>> G = (['A', 'B', 'C', 'D'], [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'B')])
        >>> w = {('A', 'B'): 1, ('B', 'C'): 3, ('B', 'D'): 1, ('C', 'D'): 8, ('D', 'B'): 2}
        >>> bellman_ford(G, w, 'D')
        ({'A': inf, 'B': 2, 'C': 5, 'D': 0},
         {'A': None, 'B': 'D', 'C': 'B', 'D': None})

    """
    d, p = initialize(G, s)
    V, E = G
    for _ in range(len(V)-1):
        for (u, v) in E:
            if d[v] > d[u] + w[u, v]:
                d[v] = d[u] + w[u, v]
                p[v] = u
    for (u, v) in E:
        if d[v] > d[u] + w[u, v]:
            raise RuntimeError('Graph contains negative cycles.')
    return d, p


def dijkstra(G, w, s):
    """Dijkstra's algorithm for shortest-path search."""
    d, p = initialize(G, s)
    V, E = G
    S = set(V)
    while S:
        u = min(S, key=lambda x: d[x])
        S = S - {u}
        for (t, v) in E:
            if t == u and d[v] > d[u] + w[u, v]:
                d[v] = d[u] + w[u, v]
                p[v] = u
    return d, p


def floyd_warshall(W):
    """Floyd-Warshall algorithm for shortest-path search between all vertexes of
    the graph W. W is represented with an adjacancy matrix.
    """
    n = len(W)
    D = {x: None for x in range(n)}
    D[0] = list(W)
    for k in range(1, n+1):
        D[k] = list(D[k-1])
        for i in range(n):
            for j in range(n):
                D[k][i][j] = min(D[k-1][i][j], D[k-1][i][k-1] + D[k-1][k-1][j])
    return D[n]


def johnson(G, w):
    """Johnson's algorithm for shortest-path search between all vertexes of
    the graph G. G is represented with an adjacancy list. The parameter w
    represents weights of edges.
    """
    V, E = G
    G_ = (V + ['S'], E + [('S', v) for v in V])
    V_, E_ = G_
    w_ = dict(w.items() + [((u, v), 0) for (u, v) in E_ if u == 'S'])
    d, p = bellman_ford(G_, w_, 'S')
    h = {}
    for v in V:
        h[v] = d[v]
    w__ = {}
    for (u, v) in E:
        w__[u, v] = w[u, v] + h[u] - h[v]
    D = {(u, v): None for u in V for v in V}
    for u in V:
        d_, p_ = dijkstra(G, w__, u)
        for v in V:
            D[u, v] = d_[v] + h[v] - h[u]
    return D
    
    
class DisjointSet(object):
    """Simple implementation of disjoint-set data structure."""
    
    def __init__(self, key):
        self.key = key
        self.parent = None
        self.rank = 0

    def __eq__(self, other):
        if isinstance(other, type(self)):
            return self.key == other.key
        return False

    def __ne__(self, other):
        return not (self == other)
        
        
def find_set(x):
    """Finds representant of the given data structure x."""
    if x.parent is None:
        return x
    x.parent = find_set(x.parent)
    return x.parent


def union(x, y):
    """Joins two subsets into a single subset."""
    x = find_set(x)
    y = find_set(y)
    if x.rank > y.rank:
        y.parent = x
    else:
        x.parent = y
        if x.rank == y.rank:
            y.rank += 1

    
def kruskal(G, w):
    """Implementation of Kruskal's algorithm. Returns minimal spanning
    tree for the given graph G and wights w.
    """
    V, E = G
    s = {}
    for v in V:
        s[v] = DisjointSet(v)
    K = []
    for (u, v) in sorted(E, key=lambda e: w[e]):
        if find_set(s[u]) != find_set(s[v]):
            K = K + [(u, v)]
            union(s[u], s[v])
    return K

    
# -------------------------------------
# Algorithms for solving other problems
# -------------------------------------

def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in xrange(0, len(l), n):
        yield l[i:i+n]


def median(s):
    """Returns median of given sequence (isn't very accurate)."""
    n = len(s)
    if n <= 5:
        return sorted(s)[n/2]
    else:
        m = [median(c) for c in chunks(s, 5)]
        return sorted(m)[int(len(m)/2)]


def select(s, i):
    """Implementation of the select algorithm."""
    m = median(s)
    l = [x for x in s if x < m]
    r = [x for x in s if x > m]
    if len(l) + 1 == i:
        return m
    if len(l) >= i:
        return select(l, i)
    else:
        return select(r, i-len(l)-1)