最近用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)
