python建立区间树

python建立区间树

此处建立的区间树和算法导论上面的略有不同，是将一个大的长为length的区间，如1到100，按照指定的分支数$k$（$k$>=2），如$k$=8划分，建立区间树，每个小区间为树上的一个节点，$k$-区间树的定义如下：
1）$k$大于等于2
2）$k$-区间树划分规则。对于每个非叶子节点 $x$：
（1）若

$\left|Seg\left(x\right)\right|\le k$，
直接划分 $\left|Seg\left(x\right)\right|$个节点
（2）若
$\left|Seg\left(x\right)\right|>k$，
划分出 $k$个孩子满足
$\lfloor \frac{|Seg\left(x\right)|}{k}\rfloor \le \left|\underset{yϵSon\left(x\right)}{Seg\left(x\right)}\right|\le \lceil \frac{|Seg\left(x\right)|}{k}\rceil$
其中 $k$为划分参数，每个非叶子节点均有 $k$个孩子的 $k$-区间树为满 $k$-区间树

import math
from functools import reduce
from queue import Queue
import numpy as np


class HNode(object):
    def __init__(self, start, end, count):
        self.start = start
        self.end = end
        self.count = count
        self.parent = None
        self.children = []
        self.height = None

    def isleaf(self):
        return len(self.children) == 0

    def __repr__(self):
        return 'HNode(start: %d, end: %d) = (count=%s)' % (
            self.start, self.end, self.count)


class HTree(object):
    # 建区间树，按照树的层次，由根到叶子节点逐步建立，思想源于二叉树的层次遍历
    def __init__(self, arity, leaf_counts):
        self.k = arity
        self.length = len(leaf_counts)
        height = int(math.ceil(math.log(len(leaf_counts), arity))) + 1
        self.height = height
        flag = [0] * len(leaf_counts)  # 设置标志位，用于最后叶子节点的判断，避免最后叶子节点无限分割下去
        q = Queue()
        root = HNode(0, len(leaf_counts) - 1, sum(leaf_counts))
        root.height = height
        self.root = root
        # last 和 nlast 用于判断当前节点的高度，思想同样源于二叉树的层次遍历
        last = root
        nlast = None
        q.put(root)
        while not q.empty():
            # 首先设一个列表，存储将某区间按照区间的个数，即树的分支数，每个小区间的长度
            # 其次，将这些小区间作为孩子节点建立父亲和孩子关系
            p = q.get()
            parent = p
            length = p.end + 1 - p.start
            l = list()
            if length <= arity:
                for i in range(length):
                    l.append(1)
            else:
                down = math.floor(length / arity)
                residue = length - down * arity
                for i in range(residue):
                    l.append(down + 1)
                for i in range(arity - residue):
                    l.append(down)
            # print("list l: " + str(l))
            begin = p.start
            # print("begin: " + str(begin))
            # 建立父亲孩子节点之间的关系，并将孩子节点加入列表，麻烦之处在于最后处理叶子节点
            # 要避免将叶子节点无限分割为start和end相等的子节点。所以设立标志位，第一次start=end
            # 时置标志位为1，之后如果在遇到start=end，直接略过，不作为子节点加入树
            for i in range(len(l)):
                node = HNode(begin, begin + l[i] - 1, sum(leaf_counts[begin:begin + l[i]]))
                if node.start != node.end:
                    node.height = height - 1
                    q.put(node)
                    nlast = node
                    # print(node)
                    self.__add_child(parent, node)
                else:
                    if flag[node.start] == 0:
                        flag[node.start] = 1
                        node.height = height - 1
                        q.put(node)
                        nlast = node
                        self.__add_child(parent, node)
                begin = begin + l[i]
            if p == last:
                last = nlast
                height = height - 1

        # 最后将虚拟下标映射到实际数据列表中，重新赋值节点的start和end
        for node in self.translevel():
            node.start = leaf_counts[node.start]
            node.end = leaf_counts[node.end]

    # def __repr__(self):
    #     return '\n'.join(map(lambda x: x.__repr__(), self.postorder_iter()))

    
    # 树的层次遍历
    def translevel(self, with_height=False):
        q = Queue()
        root = self.root
        q.put(root)
        while not q.empty():
            p = q.get()
            yield p
            for node in range(len(p.children)):
                q.put(p.children[node])
                
    # 树的后序遍历
    def postorder_iter(self):
        if self.root == None:
            return
        slist = list()
        res = list()
        ans = list()
        slist.append(self.root)
        while len(slist):
            top = slist.pop()
            res.append(top)
            for i in top.children:
                slist.append(i)
        while len(res):
            ans.append(res.pop())
        for j in ans:
            yield j

    def __add_child(self, parent, node):
        parent.children.append(node)
        node.parent = parent


def main():
    data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]
    data1 = [1, 2, 3, 4, 5, 6, 7]
    data2 = [1, 2, 3, 4, 5]
    tree = HTree(2, data2)

    for node in tree.postorder_iter():
        # if node.isleaf():
        print(node)

    # for node in tree.translevel():
    #     if node.isleaf():
    #         print(node)


if __name__ == '__main__':
    main()