Деревья поиска

Предположим, что у нас есть числовой массив, и нам требуется уметь быстро отвечать на вопрос вида «чему равна сумма всех чисел с номера a до номера b?» Наивный способ

def segment_sum(array, a, b):
    return sum(array[a:b+1])

имеет время работы, пропорциональное разности b-a. Эта разность может быть того же порядка, что и размер массива. Поэтому на практике такая ситуация оказывается малоудовлетворительной (для массива размером порядка миллиона элементов за секунду можно обработать максимум порядка тысячи запросов).

Есть классический трюк: можео заранее вычислить префиксные суммы массива

def prefix_sums(array):
    result = []

    for x in array: result.append(result[-1] + x)

    return result

Тогда сумму чисел на отрезке можно вычислить как разность двух таких сумм:

def segment_sum_prefix(prefix_sums_array, a, b):
    return prefix_sums_array[b+1] - prefix_sums_array[a]

Но сумма — далеко не единственная операция, которую приходится считать на отрезках массива. Не менее часто встречается необходимость считать, например, максимумы и минимумы. Для них вышеописанный трюк не проходит: у них нет аналога разности.

Именно для решения этой проблемы и используются деревья поиска.

Простейшее дерево

Сперва поставим задачу. Предположим, что определена ассоциативная функция combine, принимающая на вход массив и работающая за время, зависящее только от длины входного массива (но не от его содержимого).

Ассоциативность в данном случае означает, что

combine([combine(subarray) for subarray in array])

совпадает с

combine([x for subarray in array for x in subarray])

Нам нужно предобработать массив таким образом, чтобы можно было получать результат функции

def segment_value(array,a,b):
    return combine(array[a:b+1])

за время, с точностью до логарифмического множителя не зависящее от a, b или же размера массива array.

Будем для любого натурального K, большего 1, называть K-деревом результат работы следующей функции:

def make_tree(array):
    level = array[:]

    levels = [level]

    while len(level) >= K:
        level = [combine(level[i:i+K]) for i in range(0, len(level), K)]
        levels.append(level)

    return levels

Например, если определить combine = sum и K=2, то

make_tree([1,2,3,4,5,6,7,8]) == [
  [1,2,3,4,5,6,7,8],
  [3,  7,  11, 15],
  [10,     26],
  [36]
]

Теперь, наконец, мы можем решить нашу задачу при помощи следующей функции:

def tree_segment_value(tree, a, b):
    N = len(tree)

    if N == 0: return combine([])

    return tree_segment_value_rec(tree, a, b, N-1, 0, K**(N-1))


def tree_segment_value_rec(tree, a, b, level, start, size):
    to_combine = []

    for i in range(start, min(start+K, len(tree[level]))):
        if (i + 1)*size <= a: continue
        if i      *size >  b: break
        if i*size >= a and (i+1)*size-1 <= b: 
            to_combine.append(tree[level][i])
        else:
            to_combine.append(
              tree_segment_value_rec(tree, a, b, level-1, i * K, size // K)
            )
    
    return combine(to_combine)

Интересно, что эту функцию можно (хотя и не очень просто) реализовать без использования рекурсии или каких бы то ни было её замен (типа явного стека). Соответствующую реализацию оставим в качестве упражнения для читателя.

Остаётся вопрос о времени работы этой функции. Заметим следующее:

• за один цикл происходит не более двух рекуррентных вызовов
• если произошло два рекуррентных вызова, то в рамках каждого из них ни один из циклов (имеется в виду не только цикл непосредственно вызванной функции, но циклы всех её потомков) не может сделать два рекуррентных вызова
• каждый цикл имеет время работы, пропорциональное сумме K и времени работы combine на не более чем K-элементном массиве

Учитывая всё вышеперичисленное, получаем, что время работы не превосходит (помноженное на некоторую константу) удвоенное произведение количества уровней дерева и суммы K и времени работы combine на не более чем K-элементом массиве. Что не превосходит (помноженный на константу) K-ричный логарифм количества элементов исходного массива. Что и требовалось.

Дерево отрезков

Вообще говоря, подобная предобработка может быть применена не только к обычному массиву (выдающему элементы по их порядковым номерам), а к любому ассоциативному массиву (выдающему элементы по «ключам»), множество ключей которого упорядочено (т.е. ключи можно сравнивать между собой).

Единственное отличие: теперь вместе со значением combine в дереве придётся хранить ещё и отрезок — минимальное и максимальное значения ключей элементов, к которым этот combine применён.

Также следует позаботиться о том, чтобы массив был отсортирован по ключам: функция tree_segment_value (и дальнейшее рассуждение об асимптотической сложности) явно использует сортированность массива.

Конкретную реализацию дерева отрезков для ассоциативного массива оставим в качестве упражнения для читателя.

Розовые кусты

Недостатком вышеприведённой модели дерева (набор «уровней») является её статичность: она не поддерживает эффективное добавление новых элементов. Конечно, существуют относительно простые методы, позволяющие «научить» произвольную статическую структуру данных эффективному добавлению новых элементов, но они накладывают дополнительные ограничения на combine (например, динамизация, описанная в ресурсе по ссылке, требует коммутативности).

Поэтому сейчас мы рассмотрим другую модель, несколько более гибкую. А именно, деревом будем называть список деревьев, спаренный с некоторым набором данных.

Например, вышеприведённое дерево сумм для последовательности целых чисел от 1 до 8 в такой модели может выглядеть так:

(36, [
  (10, [
    (3, [
      (1, []),
      (2, []),
    ]),
    (7, [
      (3, []),
      (4, []),
    ]),
  ]),
  (26, [
    (11, [
      (5, []),
      (6, []),
    ]),
    (15, [
      (7, []),
      (8, []),
    ]),
  ]),
])

Традиционное название для такой модели — «розовый куст» (rose tree).

Очень часто встречаются специализации и модификации этой модели для конкретных размеров списков. Например, бинарное дерево обычно определяется так:

• это либо «пустое значение»
• либо пара бинарных деревьев, агрегированная с набором данных

Как нетрудно заметить, если в вышеприведённый розовый куст вставить «пустые значения»:

(36, [
  (10, [
    (3, [
      (1, [None, None]),
      (2, [None, None]),
    ]),
    (7, [
      (3, [None, None]),
      (4, [None, None]),
    ]),
  ]),
  (26, [
    (11, [
      (5, [None, None]),
      (6, [None, None]),
    ]),
    (15, [
      (7, [None, None]),
      (8, [None, None]),
    ]),
  ]),
])

то получится как раз нечто, что можно рассмативать как бинарное дерево.

Добавление элементов

Будем пользоваться следующим внутренним представлением

def tree_segment(tree):    return tree[0]
def tree_value(tree):      return tree[1]
def tree_children(tree):   return tree[2]
def make_leaf(key, value): return ( (key, key), combine([value]), [] )
def make_branch(trees):
  segments = [ tree_segment(tree) for tree in trees ]
  
  left  = min([segment[0] for segment in segments])
  right = max([segment[1] for segment in segments])
  ## при должной аккуратности можно заменить на 
  ## left  = segments[0][0]
  ## right = segments[-1][1]

  value = combine([ tree_value(tree) for tree in trees ])
  
  return ( (left, right), value, trees[:] )

def tree_segment_value(tree, a, b):
  ta, tb = tree_segment(tree)
  if a <= ta and tb <= b: return tree_value(tree)

  to_combine = []

  for child in tree_children(tree):
    ta, tb = tree_segment(child)
    if tb < a: continue
    if ta > b: break
    if a <= ta and tb <= b: to_combine.append(tree_value(child))
    else: to_combine.append(tree_segment_value(child, a, b))

  return combine(to_combine)

Для добавления элемента мы реализуем иммутабельный интерфейс (напомним, что этот термин означает, что добавление не изменяет существующее дерево, а на его основе создаёт новое):

def tree_put(tree, key, value):
  if tree == None: return make_leaf(key, value)
  
  children = tree_children(tree)

  if len(children) == 0: ## лист
    new_leaf = make_leaf(key, value)

    tree_key = tree_segment(tree)[0]
    if key < tree_key: return make_branch([new_leaf, tree])
    if key > tree_key: return make_branch([tree, new_leaf])
    return new_leaf

  i = 0
  while i < len(children):
    if tree_segment(children[i])[0] > key: break
    i += 1
  if i > 0: i -= 1
  
  return make_branch(
    children[:i] + [tree_put(children[i], key, value)] + children[i+1:]
  )

Такой способ добавления всегда создаёт бинарные деревья. У него есть два существенных недостатка:

• при неудачном стечении обстоятельств (например, добавляемые данные упорядочены по возрастанию их ключей) высота дерева пропорциональна количеству данных
• не используется потенциальная возможность веток иметь больше поддеревьев, чем два

В связи с этим мы покажем ещё один способ добавления элементов, избавленный от обоих из этих недостатков:

MAX_SIZE = 4  ## некая константа, большая 2
def is_leaf(tree): return len(tree_children(tree)) == 0

def tree_put_rec(tree, key, value):  ## возвращает одно или два дерева
  if tree == None: return [make_branch([make_leaf(key, value)])]
  
  children = tree_children(tree)

  i = 0
  while i < len(children):
    if tree_segment(children[i])[0] > key: break
    i += 1
  if i > 0: i -= 1

  left_children = children[:i]
  right_children = children[i+1:]

  child = children[i]
  if is_leaf(child):
    new_leaf = make_leaf(key, value)

    child_key = tree_segment(child)[0]
    if   child_key < key: mid_children = [child, new_leaf]
    elif child_key > key: mid_children = [new_leaf, child]
    else:                 mid_children = [new_leaf]
  else:
    mid_children = tree_put_rec(child, key, value)
  
  new_children = left_children + mid_children + right_children
  N = len(new_children)
  if N > MAX_SIZE: to_wrap = [new_children[:N//2], new_children[N//2:]]
  else:            to_wrap = [new_children]

  return [make_branch(subtrees) for subtrees in to_wrap]

def tree_put(tree, key, value):
  new_trees = tree_put_rec(tree, key, value)
  
  if len(new_trees) < 2: return new_trees[0]
  return make_branch(new_trees)

Этот способ добавления называется «B-деревом». Альтернативный способ получить дерево с логарифмической относительно количества хранимых данных высотой представлен в следующем разделе.

Для того, чтобы посмотреть вживую на эти деревья, можно воспользоваться следующим простым способом визуализации:

def print_tree(tree, indent = 0):
  if tree == None: print()

  print(' '*indent + str(tree_segment(tree)) + ' -> ' + str(tree_value(tree)))
  for child in tree_children(tree):
    print_tree(child, indent + 2)


t = None
combine = sum

for v, k in enumerate([5,7,2,1,3,6,2,7]):
  t = tree_put(t, k, v)

print_tree(t)

Балансировка Адельсон-Вельского и Ландиса

Исторически первая разновидность самобалансирующихся деревьев поиска была в 1962-м году предложена советскими математиками Г.М. Адельсон-Вельским и Е.М. Ландисом. В настоящее время на практике она почти не используется (поскольку очень сильно проигрывает в производительности алгоритмам, учитывающим наличие в современных ЭВМ большого и быстрого кэша). Здесь она присутствует исключительно из-за требований рабочей программы курса (а также — по той причине, что задача анализа высоты такого дерева представляет самостоятельный математический интерес).

Начнём с более классической (но всю эту идеологию можно применять и к построенным в предыдущем разделе бинарным «кустам отрезков») реализации бинарных деревьев: вместо того, чтобы для каждого из деревьев хранить отрезок ключей, обслуживаемый этим деревом, мы будем в каждой ветке хранить «разделяющее значение» — нечто, большее любого ключа левого поддерева, и не превышающее любого ключа правого поддерева (в этой идеологии можно продвинуться дальше и хранить данные с ключами прямо в узлах, сохраняя то же свойство — ключ в узле должен разделять ключи левого и правого поддереьвев).

Выглядит это примерно так:

## Служебные функции
def leaf(key, value): return (key, value)
def is_leaf(tree): return len(tree) == 2

def branch(separator, left, right, inv=False):
    if inv: left, right = right, left
    return (separator, left, right)

def children(tree, inv):
    if inv: return tree[2], tree[1]
    return tree[1], tree[2]


## Основной интерфейс
def empty(): return None

def put(tree, key, value):
    if tree == None: return leaf(key, value)

    tree_key = tree[0]
    inv = key < tree_key

    if is_leaf(tree):
        new_leaf = leaf(key, value)
        if key == tree_key: return new_leaf
        
        if inv: key = tree_key 
        return branch(key, tree, new_leaf, inv)
   
    left, right = children(tree, inv)
    
    return branch(tree_key, left, put(right, key, value), inv)

def get(tree, key):
    if tree == None: return None

    tree_key = tree[0]

    if is_leaf(tree):
        if tree_key == key: return tree[1]
        return None

    _, right = children(tree, key < tree_key)

    return get(right, key)

Деревья Адельсон-Вельского и Ландиса отличаются от «наивных» бинарных деревьев следующим:

• в каждом узле хранится высота дерева (длина самого длинного пути от этого узла до листа)
• в функции добавления элементов применяется дополнительная балансировка, обеспечивающая не превышающую 1 разность высот поддеревьев любого узла

Получается примерно следующее:

## Служебные функции
def leaf(key, value): return (key, value)
def is_leaf(tree): return len(tree) == 2

def branch(separator, left, right, inv=False):
    if inv: left, right = right, left
    return (separator, left, right, 1 + max(height(left), height(right)))

def children(tree, inv):
    if inv: return tree[2], tree[1]
    return tree[1], tree[2]

def height(tree):
    if tree == None: return 0
    if is_leaf(tree): return 1
    return tree[3]


## AVL-утилиты
def reassoc(tree, inv):
    l, r   = children(tree, inv)
    rl, rr = children(r, inv)
    return branch(r[0], branch(tree[0], l, rl, inv), rr, inv)

def avl(tree, inv):
    l, r = children(tree, inv)

    if height(r) - height(l) < 2: return tree

    rl, rr = children(r, inv)
    if height(rl) - height(rr) == 1: r = reassoc(r, not inv)
    
    return reassoc(branch(tree[0], l, r, inv), inv)


## Основной интерфейс
def empty(): return None

def put(tree, key, value):
    if tree == None: return leaf(key, value)

    tree_key = tree[0]
    inv = key < tree_key

    if is_leaf(tree):
        new_leaf = leaf(key, value)
        if key == tree_key: return new_leaf
        
        if inv: key = tree_key 
        return branch(key, tree, new_leaf, inv)
   
    left, right = children(tree, inv)
    
    return avl(branch(tree_key, left, put(right, key, value), inv), inv)
##         ^^^ <-- единственное отличие!

def get(tree, key):
    if tree == None: return None

    tree_key = tree[0]

    if is_leaf(tree):
        if tree_key == key: return tree[1]
        return None

    _, right = children(tree, key < tree_key)

    return get(right, key)

Отметим, что в оригинальном алгоритме предлагалось в каждом узле хранить не высоту, а разность высот поддеревьев. Но алгоритм при этом был чуть сложнее для понимания.

Анализ высоты AVL-деревьев

Пусть \(m(h)\) — минимальное количество элементов в дереве высоты \(h\).

Тогда, очевидно, \(m\) монотонно, и для \(h\geqslant2\) верно соотношение

\[ m(h) = m(h-1) + m(h-2) \]

Учитывая, что \(m(1) = 1\) и \(m(2) = 2\), получаем (согласно явной формуле для чисел Фибоначчи)

\[ m(h) = \frac{1}{\sqrt{5}}\left(\frac{1+\sqrt{5}}{2}\right)^{h} + \frac{1}{\sqrt{5}}\left(\frac{1-\sqrt{5}}{2}\right)^{h} \]

что весьма близко (отличается не более чем на 1) к первому слагаемому правой части

\[ \frac{1}{\sqrt{5}}\left(\frac{1+\sqrt{5}}{2}\right)^{h} \]

Учитывая, что при хранении \(N\) элементов выполняется неравенство

\[ N \geqslant m(h) \geqslant \frac{1}{\sqrt{5}}\left(\frac{1+\sqrt{5}}{2}\right)^{h} - 1 \]

получаем, что

\[ h \leqslant \frac{1}{\log(1+\sqrt{5})-1}\log(N\sqrt{5}) \]

Такой оценки нам вполне достаточно.

Данные в узлах

Вместо разделяющего значения можно в узлах хранить данные, ключ которых использовать как разделяющее значение.

Такой подход позволяет не обрабатывать отдельно «листья», значительно упрощая алгоритмы добавления и поиска, но усложняя алгоритм удаления.

Впрочем, удаление элементов мы и не обсуждали по банальной причине:

• обычно проще не удалять элемент физически из дерева, а всего лишь помечать его как удалённый
• если через такое дерево проходит очень большое количество элементов (то приходя в него, то уходя из него), и требуется сохранять ограниченный размер дерева, можно изредка его перестраивать: например, если количество помеченных элементов составляет половину (или какую-либо иную долю) его размера

Итак, без лишних слов (AVL-дерево в 26 строчек):

def height(tree): return 0 if tree == None else tree[4]

def branch(key, value, l, r, inv = False): 
  if inv: l, r = r, l
  return (key, value, l, r, 1 + max(height(l), height(r)))

## AVL-часть остаётся прежней
def children(tree, inv): return (tree[3], tree[2]) if inv else (tree[2], tree[3])
  
def reassoc(tree, inv):
    l, r   = children(tree, inv)
    rl, rr = children(r, inv)
    return branch(r[0], r[1], branch(tree[0], tree[1], l, rl, inv), rr, inv)

def avl(tree, inv):
    l, r = children(tree, inv)

    if height(r) - height(l) < 2: return tree

    rl, rr = children(r, inv)
    if height(rl) - height(rr) == 1: r = reassoc(r, not inv)
    
    return reassoc(branch(tree[0], tree[1], l, r, inv), inv)

## А вот само дерево поиска упрощается
def empty(): return None

def put(tree, key, value):
    if tree == None:   return branch(key, value, None, None)
    if tree[0] == key: return branch(key, value, tree[2], tree[3])

    inv = key < tree[0]
   
    left, right = children(tree, inv)
    
    return avl(branch(tree[0], tree[1], left, put(right, key, value), inv), inv)

def get(tree, key):
    if tree == None: return None
    if key == tree[0]: return tree[1]
    
    return get(children(tree, key < tree[0])[1], key)