这是解决问题的递归尝试。只有在相同深度上的所有列表都具有相同长度时,它才会起作用。否则会引发ValueError:
from collections.abc import Sequence
def get_shape(lst, shape=()):
"""
returns the shape of nested lists similarly to numpy's shape.
:param lst: the nested list
:param shape: the shape up to the current recursion depth
:return: the shape including the current depth
(finally this will be the full depth)
"""
if not isinstance(lst, Sequence):
# base case
return shape
# peek ahead and assure all lists in the next depth
# have the same length
if isinstance(lst[0], Sequence):
l = len(lst[0])
if not all(len(item) == l for item in lst):
msg = 'not all lists have the same length'
raise ValueError(msg)
shape += (len(lst), )
# recurse
shape = get_shape(lst[0], shape)
return shape
根据您的输入(以及来自 cmets 的输入),结果如下:
a1=[1,2,3,4,5,6]
b1=[[1,2,3],[4,5,6]]
print(get_shape(a1)) # (6,)
print(get_shape(b1)) # (2, 3)
print(get_shape([[0,1], [2,3,4]])) # raises ValueError
print(get_shape([[[1,2],[3,4]],[[5,6],[7,8]]])) # (2, 2, 2)
不确定最后的结果是不是你想要的。
更新
正如mkl 在 cmets 中指出的那样,上面的代码不会捕获嵌套列表的形状不一致的所有情况;例如[[0, 1], [2, [3, 4]]] 不会引发错误。
这是一个检查形状是否一致的镜头(可能有更有效的方法来做到这一点......)
from collections.abc import Sequence, Iterator
from itertools import tee, chain
def is_shape_consistent(lst: Iterator):
"""
check if all the elements of a nested list have the same
shape.
first check the 'top level' of the given lst, then flatten
it by one level and recursively check that.
:param lst:
:return:
"""
lst0, lst1 = tee(lst, 2)
try:
item0 = next(lst0)
except StopIteration:
return True
is_seq = isinstance(item0, Sequence)
if not all(is_seq == isinstance(item, Sequence) for item in lst0):
return False
if not is_seq:
return True
return is_shape_consistent(chain(*lst1))
可以这样使用:
lst0 = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
lst1 = [[0, 1, 2], [3, [4, 5]], [7, [8, 9]]]
assert is_shape_consistent(iter(lst0))
assert not is_shape_consistent(iter(lst1))