AlexanderUsatov 9 дек 2019 в 11:22

Как получить по индексу элемент из бинарного дерева за приемлемое время?

6 мин

5.3K

Привет, Хабр!

Полгода назад я задумался, как можно было бы получить элемент из бинарного дерева за O(log(N)). Ответ пришёл довольно быстро — Lazy Propagation. Но реализовать это в коде я поленился. Сейчас надо сдавать дипломный проект в университете, поэтому я занимаюсь чем угодно, только не им. Именно так я и сел это реализовывать.

Некоторые замечания:

Дерево не сбалансированное (пока что, ведь дипломный проект всё же надо написать), в следствие чего оценка O(log(N)) везде является амортизированной для случайных входных данных.
Подобной структуры данных я не нашёл, коллеги и друзья, у которых спрашивал, тоже не предложили ничего похожего. Если вам известна реализация подобной идеи, прошу сообщить мне.
В классе вершины можно было бы обойтись без поля parent, передавая родителя в методы.

Описание идеи

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

Увеличить countLefter по всему пути, который мы проходим до места вставки новой вершины
Для всех вершин, которые на одну правее вышеупомянутых, увеличить add на 1

Теперь логично ввести метод push(), который будет прибавлять add к countLefter самой вершины и обоих её потомков.

Вот каким получился класс вершины:

public class UNode {

    UNode parent;
    int key;
    int countLefter;
    int add;
    UNode left;
    UNode right;

    public UNode() {
    }

    public UNode(UNode parent, int key, int countLefter, int add, UNode left, UNode right) {
        this.parent = parent;
        this.key = key;
        this.countLefter = countLefter;
        this.add = add;
        this.left = left;
        this.right = right;
    }

    public void push() {
        countLefter += add;
        if (left != null)
            left.add += add;
        if (right != null)
            right.add += add;
        add = 0;
    }

    @Override
    public String toString() {
        return "Node{" +
                "key=" + key +
                ", countLefter=" + countLefter +
                '}';
    }
}

Отлично, теперь можно приступать к построению дерева!

Первое, что мы делаем, заходя в вершину — вызываем метод push().

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

Чтобы получить элемент по индексу, поступаем вполне очевидно: если index < countLefter текущей вершины, идём влево. Если значения равны, мы нашли вершину с заданным индексом. Иначе — идём вправо.

Удаление и добавление элемента, в принципе, не сильно отличаются от обычного бинарного дерева, за исключением изменения полей countLefter и add. Если мы после успешного добавления/удаления возвращаемся в вершину слева, то эти поля надо будет изменить. Если справа — нет.

Вот таким получился код дерева:

import java.util.LinkedList;

public class UTree {

    private UNode root;
    private int size;

    public UTree() {

    }

    public int size() {
        return size;
    }

    public int get(int index) {
        return get(index, root);
    }

    public boolean add(int key) {
        if (root == null) {
            root = new UNode(null, key, 0, 0, null, null);
            size++;
            return true;
        }
        boolean res = add(key, root);
        if (res) size++;
        return res;
    }

    public boolean remove(int key) {

        if (root == null)
            return false;
        if (key == this.root.key) {
            root.push();
            removeRoot();
            size--;
            return true;
        }

        boolean res = remove(key, root);
        if (res) size--;
        return res;
    }

    private int get(int index, UNode root) {
        root.push();
        if (index == root.countLefter)
            return root.key;
        if (index < root.countLefter)
            return get(index, root.left);
        return get(index, root.right);
    }

    private boolean add(int key, UNode root) {
        if (key == root.key) return false;
        root.push();
        if (key < root.key)
            if (root.left != null) {
                boolean res = add(key, root.left);
                if (res) {
                    root.countLefter++;
                    if (root.right != null)
                        root.right.add++;
                }
                return res;
            } else {
                root.left = new UNode(root, key, root.countLefter, 0, null, null);
                root.countLefter++;
                if (root.right != null)
                    root.right.add++;
                return true;
            }
        if (root.right != null)
            return add(key, root.right);
        else {
            root.right = new UNode(root, key, root.countLefter + 1, 0, null, null);
            return true;
        }
    }

    public boolean removeByIndex(int index) {
        if(this.root == null) return false;
        root.push();
        if (index == this.root.countLefter) {
            removeRoot();
            return true;
        }
        boolean res = removeByIndex(index, root);
        if (res) size--;
        return res;
    }

    private boolean removeByIndex(int index, UNode root) {
        if (root == null)
            return false;
        root.push();
        if (index == root.countLefter)
            return removeNode(root);
        if (index < root.countLefter) {
            boolean res = removeByIndex(index, root.left);
            if (res) {
                root.countLefter--;
                if (root.right != null)
                    root.right.add--;
            }
            return res;
        } else
            return removeByIndex(index, root.right);
    }

    private boolean removeNode(UNode root) {
        if (root.left == root.right) {
            if (root.parent.left == root)
                root.parent.left = null;
            else root.parent.right = null;
            return true;
        }
        if (root.left == null) {
            if (root.parent.left == root) {
                root.parent.left = root.right;
                root.right.add--;
            } else {
                root.parent.right = root.right;
                root.right.add--;
            }
            root.right.parent = root.parent;
            return true;
        }
        if (root.right == null) {
            if (root.parent.left == root)
                root.parent.left = root.left;
            else
                root.parent.right = root.left;
            root.left.parent = root.parent;
            return true;
        }
        UNode right = getRight(root.left);
        cut(right);
        root.key = right.key;
        root.countLefter--;
        root.right.add--;
        return true;
    }

    private boolean remove(int key, UNode root) {
        if (root == null)
            return false;
        root.push();
        if (key == root.key)
            return removeNode(root);
        if (key < root.key) {
            boolean res = remove(key, root.left);
            if (res) {
                root.countLefter--;
                if (root.right != null)
                    root.right.add--;
            }
            return res;
        } else
            return remove(key, root.right);
    }

    private void removeRoot() {
        if (root.left == root.right) {
            root = null;
            return;
        }
        if (root.left == null) {
            root = root.right;
            root.add--;
            return;
        }
        if (root.right == null) {
            root = root.left;
            return;
        }
        UNode right = getRight(root.left);
        cut(right);
        root.key = right.key;
        root.countLefter--;
        root.right.add--;
    }

    private static void cut(UNode node) {
        if (node.parent.left == node)
            node.parent.left = node.left;
        else
            node.parent.right = node.left;
        if (node.left != null)
            node.left.parent = node.parent;

    }

    private UNode getRight(UNode root) {
        root.push();
        if (root.right == null)
            return root;
        return getRight(root.right);
    }

    public void printTree() {
        printTree(root);
    }

    private void printTree(UNode root) {
        if (root == null) return;
        root.push();
        printTree(root.left);
        System.out.println(root);
        printTree(root.right);
    }

    public LinkedList<UNode> getAll(){
        LinkedList<UNode> res = new LinkedList<>();
        getAll(root, res);
        return res;
    }

    private void getAll(UNode root, LinkedList<UNode> res){
        if (root == null) return;
        root.push();
        getAll(root.left, res);
        res.add(root);
        getAll(root.right, res);

    }

}

→ Тут код можно найти на гитхабе.

Приведу некоторые результаты скорости работы. Тестирование проводилось на:

Добавление в дерево:

Итак, добавление миллиона случайных элементов в диапазоне [0; 1_000):
Около 100 мс. TreeSet справился с этой задачей примерно за 130 мс.

Добавление миллиона случайных элементов в диапазоне [0; 10_000):
Около 150 мс. TreeSet справился с этой задачей примерно за 190 мс.

Добавление миллиона случайных элементов в диапазоне [0; 100_000):
Около 320 мс. TreeSet справился с этой задачей примерно за 415 мс.

Добавление миллиона случайных элементов в диапазоне [0; 1_000_000):
Около 510 мс. TreeSet справился с этой задачей примерно за 700 мс.

Добавление миллиона случайных элементов в диапазоне [0; 10_000_000):
Около 590 мс. TreeSet справился с этой задачей примерно за 750 мс.

Теперь удаление

Случайным образом добавляем миллион чисел в дерево. Затем миллион раз пытаемся удалить случайное число. В тестах учитывается только время, затраченное на удаление.

Диапазон добавления и удаления [0; 10_000_000):
Около 740 мс. TreeSet справился с этой задачей примерно за 750 мс.

Диапазон добавления и удаления [0; 1_000_000):
Около 600 мс. TreeSet справился с этой задачей примерно за 800 мс (стало больше, чем в предыдущем тесте).

Диапазон добавления и удаления [0; 100_000):
Около 130 мс. TreeSet справился с этой задачей примерно за 160 мс.

Диапазон добавления и удаления [0; 10_000):
Около 45 мс. TreeSet справился с этой задачей примерно за 50 мс.

Диапазон добавления и удаления [0; 1_000):
Около 30 мс. TreeSet справился с этой задачей примерно за 37 мс.

Ну и, наконец, ради чего всё затевалось, доступ по индексу

У TreeSet нет такой функциональности. Так что приведу результаты только для UTree. Снова добавляем миллион элементов, потом получаем элемент по случайному индексу от 0 до количества элементов в дереве. Время учитывается только для доступа по индексу.

Диапазон добавления [0; 1000): 85 мс

Диапазон добавления [0; 10_000): 140 мс

Диапазон добавления [0; 100_000): 300 мс

Диапазон добавления [0; 1_000_000): 655 мс

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

P.S.

Озадачиться балансировкой дерева планирую после Нового года. Если я всё-таки это сделаю, здесь появится ссылка на продолжение.

Теги:

Хабы:

Как получить по индексу элемент из бинарного дерева за приемлемое время?

Публикации

Истории

Ближайшие события