forked from sympy/sympy
-
Notifications
You must be signed in to change notification settings - Fork 1
/
iterables.py
150 lines (124 loc) · 3.97 KB
/
iterables.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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def all(iterable):
"""Return True if all elements are set to True. This
function does not support predicates explicitely,
but this behaviour can be simulated easily using
list comprehension.
>>> all( [True, True, True] )
True
>>> all( [True, False, True] )
False
>>> all( [ x % 2 == 0 for x in [2, 6, 8] ] )
True
>>> all( [ x % 2 == 0 for x in [2, 6, 7] ] )
False
NOTE: Starting from Python 2.5 this a built-in.
"""
for item in iterable:
if not item:
return False
return True
def any(iterable):
"""Return True if at least one element is set to True.
This function does not support predicates explicitely,
but this behaviour can be simulated easily using
list comprehension.
>>> any( [False, False, False] )
False
>>> any( [False, True, False] )
True
>>> any( [ x % 2 == 1 for x in [2, 6, 8] ] )
False
>>> any( [ x % 2 == 1 for x in [2, 6, 7] ] )
True
NOTE: Starting from Python 2.5 this a built-in.
"""
for item in iterable:
if item:
return True
return False
def make_list(expr, kind):
"""Returns a list of elements taken from specified expresion
when it is of sequence type (Add or Mul) or singleton list
otherwise (Rational, Pow etc.).
>>> from sympy import *
>>> x, y = map(Symbol, 'xy')
>>> make_list(x*y, Mul)
[x, y]
>>> make_list(x*y, Add)
[x*y]
>>> set(make_list(x*y + y, Add)) == set([y, x*y])
True
"""
if isinstance(expr, kind):
return list(expr.args[:])
else:
return [expr]
def flatten(iterable):
"""Recursively denest iterable containers.
>>> flatten([1, 2, 3])
[1, 2, 3]
>>> flatten([1, 2, [3]])
[1, 2, 3]
>>> flatten([1, [2, 3], [4, 5]])
[1, 2, 3, 4, 5]
>>> flatten( (1,2, (1, None)) )
[1, 2, 1, None]
adapted from http://kogs-www.informatik.uni-hamburg.de/~meine/python_tricks
"""
result = []
for el in iterable:
if hasattr(el, "__iter__") and not isinstance(el, basestring):
result.extend(flatten(el))
else:
result.append(el)
return result
def postorder_traversal(node):
""" Do a postorder traversal of a tree.
This generator recursively yields nodes that it has visited in a postorder
fashion. That is, it descends through the tree depth-first to yield all of
a node's children's postorder traversal before yielding the node itself.
Parameters
----------
node : sympy expression
The expression to traverse.
Yields
------
subtree : sympy expression
All of the subtrees in the tree.
Examples
--------
>>> from sympy import symbols
>>> from sympy.utilities.iterables import postorder_traversal
>>> x,y,z = symbols('xyz')
>>> set(postorder_traversal((x+y)*z)) == set([z, y, x, x + y, z*(x + y)])
True
"""
for arg in node.args:
for subtree in postorder_traversal(arg):
yield subtree
yield node
def preorder_traversal(node):
""" Do a preorder traversal of a tree.
This generator recursively yields nodes that it has visited in a preorder
fashion. That is, it yields the current node then descends through the tree
breadth-first to yield all of a node's children's preorder traversal.
Parameters
----------
node : sympy expression
The expression to traverse.
Yields
------
subtree : sympy expression
All of the subtrees in the tree.
Examples
--------
>>> from sympy import symbols
>>> from sympy.utilities.iterables import preorder_traversal
>>> x,y,z = symbols('xyz')
>>> set(preorder_traversal((x+y)*z)) == set([z, x + y, z*(x + y), x, y])
True
"""
yield node
for arg in node.args:
for subtree in preorder_traversal(arg):
yield subtree