Как получить по индексу элемент из бинарного дерева за приемлемое время?
Привет, Хабр!
Полгода назад я задумался, как можно было бы получить элемент из бинарного дерева за 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 getAll(){
LinkedList res = new LinkedList<>();
getAll(root, res);
return res;
}
private void getAll(UNode root, LinkedList 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.
Озадачиться балансировкой дерева планирую после Нового года. Если я всё-таки это сделаю, здесь появится ссылка на продолжение.
