python实现三叉树

最近用python实现了三叉树，觉得挺有意思的，所以分享出来：

class Node:
  """The Node class. You should not change this!"""

  def __init__(self,ID,data):
    self.ID=ID
    self.data=data
    self.children = [] # This is a list of other Node objects. At first it is empty.

def getID(tree):
    return tree.ID

PRINT_NODE_LIST = []

def printNode(tree):
    print(getID(tree))
    print(getData(tree))

    PRINT_NODE_LIST.append(getID(tree))

def getData(tree):

    return tree.data

ef buildTree(data, split_list, root_id):
    """
    Build a tree using any tree data structure that you choose.

    The root node of this (sub) tree should store all of data.

    If the data is empty or the split_list is empty,
    len(data) == 0 or len(split_list) == 0,
    the this root node will have no children and it should still 
    store the id and the data (even though it is empty).

    If the data is not empty and the split_list is not empty,
    the root node of this (sub) tree should have three children.
    The children should all be tree stuctures (just like the root).
    The three children should each be given a subset of data.

    Let s = split_list[0], the first integer in split_list.
    Child index 0:
        Should contain a NumPy array with shape (N_0, M)
        that containing all points in data where data[i,s] < -1.
        N_0 is the number of points that fit this criteria.
        If N_0 is zero, this child should store an empty NumPy array
        with shape (0, M), e.g. np.zeros((0,M)).
        This child should have id = root_id*10
    Child index 1:
        The same as child index 0 except:
        Child data is -1 <= data[i,s] <= 1
        Child id = root_id*10 + 1        
    Child index 2:
        The same as child index 2 except:
        Child data is data[i,s] > 1
        Child id = root_id*10 + 2        

    The tree should continue growing where the children split their 
    data based on split_list[1] and the grandchildren split their data
    based on split_list[2] and so on.

    Input:
    data: NumPy ndarray shape (N, M) representing the M-dimensional 
        coordinates of N data points. The data may be an empty NumPy
        array, i.e. len(data) == 0.
    split_index_order: List of integers. Each integer in the list is
        in the range [0,M). The list will have at most M entries.
        The list may be empty, [], in which case, this tree will not
        have any children.
    root_id: Positive integer representing the ID for the root of this
        (sub) tree. The ID for the root any child subtrees should be
        root_id*10 + index_of_that_child. So if the root_id is 7 and 
        there are three children, the IDs of the children should be
        70, 71, 72.

    Return:
    A data structure of your choosing that represents the resulting
        tree.
    """
    ### BEGIN YOUR CODE ###
    # create the root data
    root = Node(ID=root_id,data=data)
    curr_layer = [root]
    idx=0
    if(len(data) == 0 or len(split_list) == 0):
        return root

    while(True):

        # get the split index
        index=split_list[idx]
        # store the new layer node
        temp=[]
        for node in curr_layer:
            # print('hello')
            cur_root=node
            cur_data=cur_root.data
            
            if(len(cur_data)==0):
                continue

            # select the column data
            arr=cur_data[:,index]
            child_0=[]
            child_1=[]
            child_2=[]
            for i,item in enumerate(arr):
                if(item<-1):
                    child_0.append(i)
                elif(item>=-1 and item<=1):
                    child_1.append(i)
                else:
                    child_2.append(i)
            # create children node
            for i in range(3):
                if(i==0 and len(child_0)!=0):
                    arr1=cur_data[child_0,:]
                elif(i==1 and len(child_1)!=0):
                    arr1=cur_data[child_1,:]
                elif(i==2 and len(child_2)!=0):
                    arr1=cur_data[child_2,:]
                else:
                    arr1=np.zeros((0,cur_data.shape[-1]))
                child_id=cur_root.ID*10+i
                child = Node(child_id,arr1)
                cur_root.children.append(child)
                temp.append(child)
        idx+=1
        curr_layer=[item for item in temp]
        # print('hello')
        if(idx>=len(split_list)):
            break
    return root

def printTreeBF(tree):
    list_node=[tree]
    while(len(list_node)>0):
        cur_node=list_node[0]
        list_node.pop(0)
        printNode(cur_node)
        # print(cur_node.children)
        for item in cur_node.children:
            list_node.append(item)


def printTreeDF(tree):

    printNode(tree)
    for item in tree.children:
        printTreeDF(item)

感觉还蛮有意思的，我这里仅展示核心代码哈。

测试用例也写一个：

def task1_testC(data):
    tree1C = buildTree(data, [0, 1, 2], 1)
    print("tree1C root")
    printNode(tree1C)

    tree1C_firstChild = getChildren(tree1C)[0]
    print("tree1C first child")
    printNode(tree1C_firstChild)

    tree1C_firstGrandchild = getChildren(tree1C_firstChild)[0]
    print("tree1C first grandchild")
    printNode(tree1C_firstGrandchild)
    
data1 = np.array([
        [-1.5, -0.5, -0.2],
        [0.3, 1.3, 0.0],
        [-1.3, -1.4, -2.1],
        [0.9, 1.5, -0.6]])

task1_testC(data1)