def viterbi(matrix):
    """Viterbi algorithm : find the best 8-connected path in a matrix.
    
    Example, here the best path from top to bottom is to follow the "4"s 
    (diagonal), the energy of the path is 0
    >>> viterbi([[4, 1, 3], \
                 [3, 4, 2], \
                 [2, 1, 4]])
    [0, 1, 2]
    
    On the next example, 2 best paths, only the first found is returned
    >>> viterbi([[1, 2, 2], \
                 [1, 1, 2], \
                 [1, 2, 1]])
    [0, 0, 0]
    
    A basic example. Best path energy: 1
    >>> viterbi([[1, 0], \
                 [4, 2]])
    [0, 1]
    """
    # init
    h = len(matrix) - 1
    w = len(matrix[0])
    k = 0
    S = {}      #Survivors
    L = {}      #Accumulate Length
    for i in range(w):
        L[(k, i)] = 0

    # recursion
    for k in range(h):
        for x in range(w):
            for dx in range(-1, 2):
                if (x + dx) < 0 or (x + dx) >= w:
                    continue
                tmpl = L[k, x] + abs(matrix[k][x] - matrix[k+1][x+dx])
                if (not (k+1, x+dx) in L) or (tmpl < L[(k+1, x+dx)]):
                    L[(k+1, x+dx)] = tmpl
                    S[(k+1, x+dx)] = x

    # look for the smallest
    tmp, idx = -1, -1
    for i in range(w):
        if L[(h, i)] < tmp or tmp == -1:
            tmp = L[(h, i)]
            idx = i

    # create the inversed path
    res = [idx]
    for k in range(h, 0, -1):
        idx = S[(k, idx)]
        res.append(idx)

    # reverse the path
    res.reverse()
    return res

if __name__ == "__main__":
    import doctest
    doctest.testmod()
