algorithm-in-python/dataStructure/intervalTree.py

''' mbinary
#########################################################################
# File : intervalTree.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.xyz
# Github: https://github.com/mbinary
# Created Time: 2018-11-18  10:48
# Description:
#########################################################################
'''
from random import randint, shuffle
from redBlackTree import redBlackTree

from functools import total_ordering


@total_ordering
class node:
    def __init__(self, low, high, left=None, right=None, isBlack=False):
        self.val = low   # self.val is the low
        self.high = high
        self.max = high
        self.left = left
        self.right = right
        self.parent = None
        self.isBlack = isBlack

    def __lt__(self, nd):
        return self.val < nd.val

    def __eq__(self, nd):
        return nd is not None and self.val == nd.val

    def setChild(self, nd, isLeft=True):
        if isLeft:
            self.left = nd
        else:
            self.right = nd
        if nd is not None:
            nd.parent = self

    def getChild(self, isLeft):
        if isLeft:
            return self.left
        else:
            return self.right

    def __bool__(self):
        return self.val is not None

    def __str__(self):
        color = 'B' if self.isBlack else 'R'
        return f'{color}[{self.val},{self.high}]-{self.max}'

    def __repr__(self):
        return f'intervalNode({self.val},{self.high},{self.max},isBlack={self.isBlack})'

    def overlap(self, low, high):
        return self.val <= high and self.high >= low

    def setMax(self):
        l = 0 if self.left is None else self.left.max
        r = 0 if self.right is None else self.right.max
        self.max = max(self.high, l, r)
        return self.max


class intervalTree(redBlackTree):
    def search(self, low, high):
        nd = self.root
        while nd is not None and not nd.overlap(low, high):
            if nd.left is not None and nd.left.max >= low:
                nd = nd.left
            else:
                nd = nd.right
        return nd

    def insert(self, nd):
        super(intervalTree, self).insert(nd)
        while nd is not None:
            nd.setMax()
            nd = nd.parent

    def delete(self, val):
        nd = self.find(val)
        if nd is not None:
            nd.max = 0
            tmp = nd.parent
            while tmp is not None:
                tmp.setMax()
                tmp = tmp.parent
            super(intervalTree, self).delete(val)

    def rotate(self, prt, chd):
        '''rotate prt, and return new prt, namyly the original chd'''
        super(intervalTree, self).rotate(prt, chd)
        prt.setMax()
        chd.setMax()

    def copyNode(self, src, des):
        des.val = src.val
        des.high = src.high
        des.setMax()


def genNum(n=10, upper=10):
    nums = {}
    for i in range(n):
        while 1:
            d = randint(0, 100)
            if d not in nums:
                nums[d] = (d, randint(d, d+upper))
                break
    return nums.values()


def buildTree(n=10, nums=None, visitor=None):
    #if nums is None or nums ==[]: nums = genNum(n)
    tree = intervalTree()
    print(f'build a red-black tree using {nums}')
    for i in nums:
        tree.insert(node(*i))
        if visitor:
            visitor(tree, i)
    return tree, nums


def testInsert(nums=None):
    def visitor(t, val):
        print('inserting', val)
        print(t)
    tree, nums = buildTree(visitor=visitor, nums=nums)
    print('-'*5 + 'in-order visit' + '-'*5)
    for i, j in enumerate(tree.sort()):
        print(f'{i+1}: {j}')
    return tree


def testSuc(nums=None):
    tree, nums = buildTree(nums=nums)
    for i in tree.sort():
        print(f'{i}\'s suc is {tree.getSuccessor(i)}')


def testDelete(nums=None):
    tree, nums = buildTree(nums=nums)
    print(tree)
    for i in nums:
        print(f'deleting {i}')
        tree.delete(i[0])
        print(tree)
    return tree


if __name__ == '__main__':
    lst = [(0, 3), (5, 8), (6, 10), (26, 26), (25, 30),
           (8, 9), (19, 20), (15, 23), (16, 21), (17, 19)]
    #lst = None
    # testSuc(lst)
    tree = testInsert(lst)
    #tree,_= buildTree(lst)
    while 1:
        a = int(input('low:'))
        b = int(input('high:'))
        res = tree.search(a, b)
        print(res)