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

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

    Полгода назад я задумался, как можно было бы получить элемент из бинарного дерева за 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.

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

    Подробнее
    Реклама

    Комментарии 17

      +1
      Я точно не помню название этой структуры данных, но по сути вы переизобрели augmented binary tree, хранящее в каждом узле количество дочерних узлов (вес). Такая структура данных позволяет производить поиск/вставку/удаление узлов по индексу за O(log(N)), однако теряется возможность быстрого поиска/вставки/удаления по значению (сами данные).

      Алгоритм поиска узла в таком дереве простой:
      — Если искомый индекс меньше веса левого поддерева, переходим к левому дочернему узлу.
      — Если искомый индекс больше или равен весу левого поддерева + 1, вычитаем из нашего индекса вес левого поддерева + 1 и переходим к правому дочернему узлу.
      — Иначе мы нашли наш узел.

      Коррексия весов производится на этапе ребалансировки дерева, после операций вставки/удаления.
        +1
        Спасибо за комментарий, но у меня не теряется возможность быстрого поиска/вставки/удаления по значению. Можете посмотреть, в статье приведены результаты скорости выполнения для добавления и удаления по ключу. Всё производится за O(log(N)) как по ключу, так и по индексу
          0
          Есть более простой и элегантный алгоритм(на мой взгляд):
          Хранить в каждой ноде размер ее поддерева(считая ее). Назовем cnt.
          Как осуществляется поиск:
          ищем элемент с номером X:
          если у левого поддерева этот cnt >=X — идем туда, если cnt левого= X-1, то это сама нода, где мы находимся, иначе идем искать в правом поддереве, уменьшив X на cnt левого +1
          +1

          А почему теряются операции по значению? Кажется, если по индексу разрешить только поиск и удаление, то поиск/вставку/удаление по значению также быстро можно оставить.

          +1
          1. Идея похожа на nested set, но более эффективная при изменении дерева.
          2. IMHO, если вам нужно обращаться к узлу дерева по его индексу, то что-то где-то пошло не так, и возможно есть более оптимальный способ хранения данных.
            +1
            Советую обратить внимание на Декартово дерево (treap, cartesian tree). В разделе про «неявные декартовы деревья» описана реализация массива со вставкой в середину за O(log(n)). В том варианте, который там предложен, нет возможности доступа по ключу, только по индексу, но не думаю, что это сложно будет сделать.

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

            P. S. Чтобы понять, о чем вообще эта статья, нужно заглянуть под кат. Не надо так писать. Публикация — это не просто изложение потока мыслей.

              0
              в неявном трипе невозможно искать (а не просто сложно реализовать) по значению (в лучшем случае можно просто рядом положить какое-то бинарное дерево по значениям для проверки, но индекс вроде как не получить), т.к. они не упорядочены и обновляются.
                0
                Наверно я не с той стороны зашел. Тут уже достаточно много написали про то, что исходная задача (быстрый поиск порядковой статистики) вроде как решается добавлением дополнительного поля в узел, которое отслеживает суммарный размер поддерева с корнем в этом узле, в том числе и обычное декартово дерево. Неявное декартово дерево я больше привел в качестве примера, где эта величина вычисляется и используется.
              +1
              Не совсем понятно, зачем 2 каунтера, достаточно одного — сколько элементов слева. Ищем индекс N, у вершины слева Z элементов. При N = Z мы нашли нужный, при N > Z идём вправо с индексом N — Z — 1, при N < Z идём влево с индексом N.
                +1
                Каунтер только один — countLefter. Поле add нужно для Lazy Propagation, чтобы не менять каунтеры у всех вершин, которые правее
                  +1
                  Так в варианте, о котором я говорю, нужно менять каунтер только у вершин на пути, по которому мы нашли элемент и только если мы при этом шли влево. Это же происходит и у вас, только ещё иногда и с полем add. Единственный «сложный» случай — удаление, нужно будет менять ещё и у какой-то из ветвей.
                0
                Подобной структуры данных я не нашёл, коллеги и друзья, у которых спрашивал, тоже не предложили ничего похожего. Если вам известна реализация подобной идеи, прошу сообщить мне.
                Вы на верном пути. Еще несколько шагов и вы изобретёте колесо AVL дерево.
                Где для узла храниться только указатели на левы и правый узел и перекос, который поддерживается в диапазоне от -1 до +1.
                  0

                  Про АВЛ-дерево я знаю)

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

                    Плюсы: ничего не надо придумывать.
                    Минусы: не достойно статьи на Хабре.

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

                    Самое читаемое