Kd-tree_and_R-tree_realisation/rtree.py

from extent import *
from math import ceil

class Entry:
    def __init__(self, extent = None, child = None, parent = None, node = None, id = None, geometry = None):
        self.MBR = extent
        self.child = child
        self.node = node
        self.id = id
        self.geometry = geometry

    def __repr__(self):
        return str(self.MBR)
    

class RtreeNode:
    def __init__(self, M, parent = None):
        self.entries = []
        self.M = M
        self.parent = parent
        self.extent = None

    def __getitem__(self, i):
        if i >= self.M or i >= len(self.entries): return None

        return self.entries[i]

    def __repr__(self):
        return str(self.extent)
    # проверка на лист
    def is_leaf(self):
        for e in self.entries:
            if e.child is not None:
                return False
        return True
    # проверка на родителя
    def is_root(self):
        return self.parent is None
    # обновление MBR узла, чтобы он включал в себе все его объекты
    def update(self):
        if not len(self.entries):
            return

        if self.entries[0] is not None:
            self.extent = self.entries[0].MBR
        
        for e in self.entries[1:]:
            self.extent = union_extent(self.extent, e.MBR)
    # обновление всего дерева до корня
    def update_up(self):
        self.update()

        if self.is_root(): 
            return

        self.parent.MBR = self.extent
        self.parent.node.update_up()
# получение MBR объектов в виде объекта класса Extent
def union_extent(e1, e2):
    xmin = min(e1.xmin, e2.xmin)
    xmax = max(e1.xmax, e2.xmax)
    ymin = min(e1.ymin, e2.ymin)
    ymax = max(e1.ymax, e2.ymax)

    return Extent(xmin, xmax, ymin, ymax)
# функция вставки
def insert(node, e, child = None):
    for ent in node.entries: # если объект уже есть в узле, то возвращаем True
        if ent.MBR == e:
            return True
    # создание объекта класса Entry с переданными параметрами
    entry = Entry(extent = e, child = child, id = e.id, geometry = e.geometry)

    if len(node.entries) < node.M: # если в узле ещё есть место, то просто добавляем объект
        entry.node = node

        if entry.child is not None: # если у объекта есть потомки, то устанавливаем им объект в качестве родителя
            entry.child.parent = entry
        
        node.entries.append(entry) # добавляем объект
        node.update_up() # обновляем узел
        
        return True
    # если места нет, то необходимо разделить узел
    M = node.M # максимально возможное количество объектов
    m = ceil(float(M) / 2) # возможное количество объектов в новых узлах
    L1 = RtreeNode(M) # новые узлы
    L2 = RtreeNode(M)
    max_i, max_j = -1, -1 # переменные для получения 2 наиболее удаленных друг от друга MBR
    max_distance = 0.0 # максимальное расстояние между MBR
    tmp_entries = [ent for ent in node.entries] # временное хранилище объектов
    tmp_entries.append(entry) # добавление нового объекта во временное хранилище
    M1 = len(tmp_entries) # размер временного хранилища объектов
    # поиск 2 наиболее удаленных друг от друга MBR
    for i in range(M1):
        for j in range(i+1, M1):
            distance = tmp_entries[i].MBR.get_distance(tmp_entries[j].MBR)
            if distance > max_distance:
                max_distance = distance
                max_i = i
                max_j = j
    # запись 2 наиболее отдалённых друг от друга MBR
    e1 = tmp_entries[max_i]
    e2 = tmp_entries[max_j]
    all_exts = [] # массив для всех оставшихся объектов
    # добавление оставшихся объектов
    for ext in tmp_entries:
        if ext is not e1 and ext is not e2:
            all_exts.append(ext)
    # добавление в новые узлы полученных объектов и последующее обновление узлов
    L1.entries.append(e1)
    L2.entries.append(e2)
    L1.update()
    L2.update()
    # мы продолжаем вставлять остальные MBR в каждый из новых узлов, так чтобы увеличение общей площади MBR в каждом узле было минимально
    while len(all_exts):
        num_remained = len(all_exts)
        current_node = None

        if len(L1.entries) == (m - num_remained):
            current_node = L1
        elif len(L2.entries) == (m - num_remained):
            current_node = L2
        
        if current_node is not None:
            while len(all_exts): 
                ext = all_exts.pop()
                current_node.entries.append(ext)
        else: 
            minimal_area = union_extent(L1.extent, L2.extent).get_area()
            minimal_next = -1
            current_node = None

            for i in range(len(all_exts)):
                tmp_ext1 = union_extent(L1.extent, all_exts[i].MBR)
                tmp_area1 = tmp_ext1.get_area() - L1.extent.get_area()
                tmp_ext2 = union_extent(L2.extent, all_exts[i].MBR)
                tmp_area2 = tmp_ext2.get_area() - L2.extent.get_area()

                if min(tmp_area1, tmp_area2) > minimal_area:
                    continue
                
                minimal_next = i

                if tmp_area1 < tmp_area2 and tmp_area1 < minimal_area:
                    tmp_current_node = L1
                    minimal_area = tmp_area1
                elif tmp_area1 > tmp_area2 and tmp_area2 < minimal_area:
                    tmp_current_node = L2
                    minimal_area = tmp_area2
                else:
                    minimal_area = tmp_area1

                    if L1.extent.get_area() < L2.extent.get_area():
                        tmp_current_node = L1
                    elif L1.extent.get_area() > L2.extent.get_area():
                        tmp_current_node = L2
                    else:
                        if len(L1.entries) < len(L2.entries):
                            tmp_current_node = L1
                        else:
                            tmp_current_node = L2
            
            if minimal_next != -1 and tmp_current_node is not None:
                ext = all_exts.pop(minimal_next)
                current_node = tmp_current_node
            
            current_node.entries.append(ext)
        
        current_node.update()

        for ent in L1.entries:
            ent.node = L1

            if ent.child is not None:
                ent.child.parent = ent
    
    # после вставки всех объектов
    split(node, L1, L2) # разделение узла
    # обновление полученных узлов
    L1.update_up()
    L2.update_up()

    return True


def split(node, L1, L2):
    # получаем MBR узлов
    entry1 = Entry(L1.extent)
    entry2 = Entry(L2.extent)
    # если расщепляется корень, то удаляем текущие записи в корне и повторно заполняем его новыми записями
    if node.is_root():
        node.entries = [] 
        entry1.node = node
        entry2.node = node
        entry1.child = L1
        entry2.child = L2

        node.entries.append(entry1)
        node.entries.append(entry2)

        L1.parent = entry1
        L2.parent = entry2
        return
    else: # иначе один из узлов заменяется на новый, а другой вставляется в родительский узел через insert
        entry1.node = L1
        L1.parent = node.parent
        L1.parent.child = L1
        
        del node
        
        insert(L1.parent.node, L2.extent, L2)
        return


def search_rtree_extent(node, e): # поиск по rtree через заданный MBR
    if node.is_leaf(): # если текущий узел - лист, то возвращаем его
        return node
    
    best_entry = None # MBR с наибольшим пересечением к заданному
    intersect_area = -1 # максимальная зона покрытия

    for ent in node.entries:
        tmp_area = ent.MBR.get_intersect(e) # зона покрытия между заданным MBR и текущим
        if tmp_area > intersect_area: # если полученная зона больше максимальной, то изменяем MBR с наибольшим покрытием и максимальную зону
            intersect_area = tmp_area 
            best_entry = ent
        
    return search_rtree_extent(best_entry.child, e) # рекурсивно идём в следующий узел, который является потомком объекта с наибольшем MBR