nested methods
Why do you nest the get_max_shape
etcetera in the pad
? There is no need to do this.
get_max_shape
Here you use recursion and a global variable. A simpler way would be to have a generator that recursively runs through the array, and yields the level and length of that part, and then another function to aggregate this results. That way you can avoid passing
def get_dimensions(array, level=0):
yield level, len(array)
try:
for row in array:
yield from get_dimensions(row, level + 1)
except TypeError: #not an iterable
pass
[(0, 3), (1, 3), (2, 3), (2, 1), (2, 2), (1, 1), (2, 1), (1, 2), (2, 2), (2, 1)]
The aggregation can be very simple using collections.defaultdict
:
def get_max_dimensions(array):
dimensions = defaultdict(int)
for level, length in get_dimensions(array):
dimensions[level] = max(dimensions[level], length)
return [value for _, value in sorted(dimensions.items())]
[3, 3, 3]
creating the result
Instead of r = np.zeros(dims, dtype=dtype) + pad_value
you can use np.full
You iterate over all possible indices, and check whether it is present in the original array. Depening on how "full" the original array is, this can save some time. It also allows you to do this without your custom get_item
method to get the element at the nested index
def iterate_nested_array(array, index=()):
try:
for idx, row in enumerate(array):
yield from iterate_nested_array(row, (*index, idx))
except TypeError: # final level
for idx, item in enumerate(array):
yield (*index, idx), item
[((0, 0, 0), 0), ((0, 0, 1), 1), ((0, 0, 2), 2), ((0, 1, 0), 3), ((0, 2, 0), 4), ((0, 2, 1), 5), ((1, 0, 0), 6), ((2, 0, 0), 7), ((2, 0, 1), 8), ((2, 1, 0), 9)]
def pad(array, fill_value):
dimensions = get_max_shape(test_array)
result = np.full(dimensions, fill_value)
for index, value in iterate_nested_array(test_array):
result[index] = value
return result
array([[[ 0, 1, 2], [ 3, -1, -1], [ 4, 5, -1]], [[ 6, -1, -1], [-1, -1, -1], [-1, -1, -1]], [[ 7, 8, -1], [ 9, -1, -1], [-1, -1, -1]]])