-
Notifications
You must be signed in to change notification settings - Fork 0
/
TreeUtils.py
102 lines (88 loc) · 3.04 KB
/
TreeUtils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
class Tree(object):
def __init__(self, index):
self.parent = None
self.is_left = False
self.index = index
self.left_children = list()
self.left_num = 0
self.right_children = list()
self.right_num = 0
self._depth = -1
self.order = []
self._e = None
def isleaf(self):
return self.left_num == 0 and self.right_num == 0
def add_left(self, child):
"""
:param child: a Tree object represent the child
:return:
"""
child.parent = self
child.is_left = True
self.left_children.append(child)
self.left_num += 1
def add_right(self, child):
"""
:param child: a Tree object represent the child
:return:
"""
child.parent = self
child.is_left = False
self.right_children.append(child)
self.right_num += 1
def size(self): # compute the total size of the Tree
if hasattr(self, '_size'):
return self._size
count = 1
for i in range(self.left_num):
count += self.left_children[i].size()
for i in range(self.right_num):
count += self.right_children[i].size()
self._size = count
return self._size
def depth(self): # compute the depth of the Tree
if self._depth > 0:
return self._depth
count = 0
if self.left_num + self.right_num > 0:
for i in range(self.left_num):
child_depth = self.left_children[i].depth()
if child_depth > count:
count = child_depth
for i in range(self.right_num):
child_depth = self.right_children[i].depth()
if child_depth > count:
count = child_depth
count += 1
self._depth = count
return self._depth
def traverse(self): # traverse the Tree
if len(self.order) > 0:
return self.order
for i in range(self.left_num):
left_order = self.left_children[i].traverse()
self.order.extend(left_order)
for i in range(self.right_num):
right_order = self.right_children[i].traverse()
self.order.extend(right_order)
self.order.append(self.index) # append the root
return self.order
def creatTree(heads):
tree = []
# current sentence has already been numberized [form, head, rel]
root = None
for idx, head in enumerate(heads):
tree.append(Tree(idx))
for idx, head in enumerate(heads):
if head == -1: # -1 mszhang, 0 kiro
root = tree[idx]
continue
if head < 0:
print('error: multi roots')
if head > idx:
tree[head].add_left(tree[idx])
if head < idx:
tree[head].add_right(tree[idx])
if head == idx:
print('error: head is it self.')
return root, tree