suffix array sort

suffix_array_sort.py

def trivial_suffix_sort(txt):
    n = len(txt)
    txt2 = txt + txt
    tmp = sorted((txt2[i:i+n], i) for i in xrange(n))
    return [x[1] for x in tmp]

def suffix_sort(txt):
    # рассматривает txt как кольцо, и сортирует строки txt[i:i + len(l)],
    # возвращает список индексов sort_permutation[i], где начинается i-ая сортированая строка
    # см. также эквивалентную функцию trivial_suffix_sort
    # для сортировки именно как суффиксов можно приклеить "самый маленький символ" в конец
    # чтобы разобраться в этом коде, нужно хорошо понимать алгоритм radix sort
    n = len(txt)
    group_start = [0] * max(n, 256)
    sort_permutation = [0] * n
    group = [0] *n
    ptmp = [0] * n
    group_tmp = [0] * n

    # сортируем подстроки по первому символу
    for ch in txt:
        group_start[ord(ch)] += 1
    prev = 0
    for i in xrange(len(group_start)):
        #конвертируем количество символов в индекс начала записи
        group_start[i], prev = prev, prev + group_start[i]
    group_start_saved = list(group_start) #ещё понадобится

    for i, ch in enumerate(txt):
        j = group_start[ord(ch)]
        sort_permutation[j] = i
        group[i] = ord(ch)
        group_start[ord(ch)] += 1

    group_start = group_start_saved

    #инвариант цикла:
    #sort_permutation - сортированные подстроки длины l (i-тая строка в сортированом массиве - txt[sort_permutation[i]:sort_permutation[i]+l])
    #group - класс эквивалентности этих подстрок,
    #    txt[i:i+l] меньше/больше/равно txt[j:j+l] <=> group[i]  меньше/больше/равно group[j],
    #    формируются как номер группы одинаковых подстрок (но на первой итерации - просто номер символа)
    #group_start - в какое место скопировать подстроку этого класса при radix sort

    l = 1
    while l < n:
        for i in xrange(n):
            #при сортировке подстрок длины 2*l по 'младшим' вторым половинкам длины l - достаточно сместить уже сортированый массив на l
            j = (sort_permutation[i] - l) % n
            cl = group[j]
            ptmp[group_start[cl]] = j #radix sort: то что отсортировано по вторым половинкам - теперь сортируем по первым 'старшим'
            group_start[cl] += 1

        sort_permutation, ptmp = ptmp, sort_permutation

        new_class = -1
        prev_cl = None
        #проходим по подстрокам длины 2*l в порядке возрастания, для нумерации групп одинаковых строк
        for i in xrange(n):
            j = sort_permutation[i]
            cl = (group[j], group[(j+l) % n])
            if cl != prev_cl:
                new_class += 1
                group_start[new_class] = i
            group_tmp[j] = new_class
            prev_cl = cl
        if new_class == n-1:
            #все подстроки разные, можно не продолжать
            break
        group, group_tmp = group_tmp, group
        l *= 2

    return sort_permutation



def test():
    def check_str(txt):
        txt += '~' #last char - to avoid ambigous suffix sort
        sa = suffix_sort(txt)
        sa2 = trivial_suffix_sort(txt)
        assert sa == sa2
    import itertools
    for k in xrange(10):
        for t in itertools.product('abc', repeat=k):
            txt = ''.join(t)
            check_str(txt)

    txt = 'three sweet switched swiss witches watch three washed swiss witch swatch watch switches. which sweet switched swiss witch watches which washed swiss witch swatch watch switch?'
    check_str(txt)
    for i in suffix_sort(txt):
        print txt[i:min(len(txt), i + 30)]

if __name__ == "__main__":
    test()