1
2
3
4 """ Implements a class used to represent N-ary trees
5
6 """
7 from __future__ import print_function
8
9 from rdkit.six.moves import cPickle
10
11
12
13
14
15
17 """ This is your bog standard Tree class.
18
19 the root of the tree is just a TreeNode like all other members.
20 """
21
22 - def __init__(self, parent, name, label=None, data=None, level=0, isTerminal=0):
23 """ constructor
24
25 **Arguments**
26
27 - parent: the parent of this node in the tree
28
29 - name: the name of the node
30
31 - label: the node's label (should be an integer)
32
33 - data: an optional data field
34
35 - level: an integer indicating the level of this node in the hierarchy
36 (used for printing)
37
38 - isTerminal: flags a node as being terminal. This is useful for those
39 times when it's useful to know such things.
40
41 """
42 self.children = []
43 self.parent = parent
44 self.name = name
45 self.data = data
46 self.terminalNode = isTerminal
47 self.label = label
48 self.level = level
49 self.examples = []
50
52 """ Set the names of each node in the tree from a list of variable names.
53
54 **Arguments**
55
56 - varNames: a list of names to be assigned
57
58 **Notes**
59
60 1) this works its magic by recursively traversing all children
61
62 2) The assumption is made here that the varNames list can be indexed
63 by the labels of tree nodes
64
65 """
66 if self.GetTerminal():
67 return
68 else:
69 for child in self.GetChildren():
70 child.NameTree(varNames)
71 self.SetName(varNames[self.GetLabel()])
72
73 NameModel = NameTree
74
76 """ Adds a TreeNode to the local list of children
77
78 **Arguments**
79
80 - node: the node to be added
81
82 **Note**
83
84 the level of the node (used in printing) is set as well
85
86 """
87 node.SetLevel(self.level + 1)
88 self.children.append(node)
89
90 - def AddChild(self, name, label=None, data=None, isTerminal=0):
91 """ Creates a new TreeNode and adds a child to the tree
92
93 **Arguments**
94
95 - name: the name of the new node
96
97 - label: the label of the new node (should be an integer)
98
99 - data: the data to be stored in the new node
100
101 - isTerminal: a toggle to indicate whether or not the new node is
102 a terminal (leaf) node.
103
104 **Returns*
105
106 the _TreeNode_ which is constructed
107
108 """
109 child = TreeNode(self, name, label, data, level=self.level + 1, isTerminal=isTerminal)
110 self.children.append(child)
111 return child
112
114 """ Removes the child node
115
116 **Arguments**
117
118 - child: a TreeNode
119
120 """
121 self.children.remove(child)
122
124 """ Replaces a given child with a new one
125
126 **Arguments**
127
128 - index: an integer
129
130 - child: a TreeNode
131
132 """
133 self.children[index] = newChild
134
136 """ Returns a python list of the children of this node
137
138 """
139 return self.children
140
142 """ Destroys this node and all of its children
143
144 """
145 for child in self.children:
146 child.Destroy()
147 self.children = []
148
149 self.parent = None
150
152 """ Returns the name of this node
153
154 """
155 return self.name
156
158 """ Sets the name of this node
159
160 """
161 self.name = name
162
164 """ Returns the data stored at this node
165
166 """
167 return self.data
168
170 """ Sets the data stored at this node
171
172 """
173 self.data = data
174
176 """ Returns whether or not this node is terminal
177
178 """
179 return self.terminalNode
180
182 """ Sets whether or not this node is terminal
183
184 """
185 self.terminalNode = isTerminal
186
188 """ Returns the label of this node
189
190 """
191 return self.label
192
194 """ Sets the label of this node (should be an integer)
195
196 """
197 self.label = label
198
200 """ Returns the level of this node
201
202 """
203 return self.level
204
206 """ Sets the level of this node
207
208 """
209 self.level = level
210
212 """ Returns the parent of this node
213
214 """
215 return self.parent
216
218 """ Sets the parent of this node
219
220 """
221 self.parent = parent
222
223 - def Print(self, level=0, showData=0):
224 """ Pretty prints the tree
225
226 **Arguments**
227
228 - level: sets the number of spaces to be added at the beginning of the output
229
230 - showData: if this is nonzero, the node's _data_ value will be printed as well
231
232 **Note**
233
234 this works recursively
235
236 """
237 if showData:
238 print('%s%s: %s' % (' ' * level, self.name, str(self.data)))
239 else:
240 print('%s%s' % (' ' * level, self.name))
241
242 for child in self.children:
243 child.Print(level + 1, showData=showData)
244
245 - def Pickle(self, fileName='foo.pkl'):
246 """ Pickles the tree and writes it to disk
247
248 """
249 with open(fileName, 'wb+') as pFile:
250 cPickle.dump(self, pFile)
251
253 """ returns a string representation of the tree
254
255 **Note**
256
257 this works recursively
258
259 """
260 here = '%s%s\n' % (' ' * self.level, self.name)
261 for child in self.children:
262 here = here + str(child)
263 return here
264
266 """ allows tree1 == tree2
267
268 **Note**
269
270 This works recursively
271 """
272 return (self < other) * -1 or (other < self) * 1
273
275 """ allows tree1 < tree2
276
277 **Note**
278
279 This works recursively
280 """
281 try:
282 nChildren = len(self.children)
283 oChildren = len(other.children)
284 if str(type(self)) < str(type(other)):
285 return True
286 if self.name < other.name:
287 return True
288 if self.label is not None:
289 if other.label is not None:
290 if self.label < other.label:
291 return True
292 else:
293 return False
294 elif other.label is not None:
295 return True
296 if nChildren < oChildren:
297 return True
298 if nChildren > oChildren:
299 return False
300 for i in range(nChildren):
301 if self.children[i] < other.children[i]:
302 return True
303 except AttributeError:
304 return True
305 return False
306
308 return not self < other and not other < self
309
310
312 tree = TreeNode(None, 'root')
313 for i in range(3):
314 tree.AddChild('child %d' % i)
315 print(tree)
316 tree.GetChildren()[1].AddChild('grandchild')
317 tree.GetChildren()[1].AddChild('grandchild2')
318 tree.GetChildren()[1].AddChild('grandchild3')
319 print(tree)
320 tree.Pickle('save.pkl')
321 print('prune')
322 tree.PruneChild(tree.GetChildren()[1])
323 print('done')
324 print(tree)
325
326 import copy
327 tree2 = copy.deepcopy(tree)
328 print('tree==tree2', tree == tree2)
329
330 foo = [tree]
331 print('tree in [tree]:', tree in foo, foo.index(tree))
332 print('tree2 in [tree]:', tree2 in foo, foo.index(tree2))
333
334 tree2.GetChildren()[1].AddChild('grandchild4')
335 print('tree==tree2', tree == tree2)
336 tree.Destroy()
337
338
339 if __name__ == '__main__':
340 _exampleCode()
341