用计数来替换嵌套列表中的元素

Question

I have a nested list:

x=[[[0, 1, 2], [0, 1, 2]], [[0, 1], [0, 1, 2]], [[0]], [[0, 1], [0, 1, 2, 3, 4]]]

Now, the goal is to get a nested list with the same structure but with elements replaced by their "global" counting numbers. So, the desired output should look like this:

y=[[[0, 1, 2], [3, 4, 5]], [[6, 7], [8, 9, 10]], [[11]], [[12, 13], [14, 15, 16, 17, 18]]]

I fight with it for the last couple of hours but without success.

Ideally, I'd like to have a universal solution being able to work with an arbitrary depth of nesting.

Any help would be very much appreciated. Thank you in advance!

Answer 1

Here's a recursive solution that does the replacement in-place and relies on the type of the element being replaced. The idea is to keep track of the "global counter" and pass it into the recursive calls so that it knows what to replace elements with:

x = [[[0, 1, 2], [0, 1, 2]], [[0, 1], [0, 1, 2]], [[0]], [[0, 1], [0, 1, 2, 3, 4]]]

def replace(lst, i):
    for j in range(len(lst)):
        if isinstance(lst[j], list):
            lst[j], i = replace(lst[j], i)
        else:
            lst[j] = i
        i += 1
    return lst, i - 1

replace(x, 0)
print(x)
# [[[0, 1, 2], [3, 4, 5]], [[6, 7], [8, 9, 10]], [[11]], [[12, 13], [14, 15, 16, 17, 18]]]

Answer 2

Here's another recursive solution. It uses itertools.count and builds a new list. Personally, I like to avoid integer indexing when possible for readability.

from itertools import count

def structured_enumerate(lst, counter=None):
    'enumerate elements in nested list, preserve structure'
    result = []
    if counter is None:
        counter = count()

    for x in lst:
        if isinstance(x, list):
            result.append(structured_enumerate(x, counter))
        else:
            result.append(next(counter))
    return result

Demo:

>>> x = [[[0, 1, 2], [0, 1, 2]], [[0, 1], [0, 1, 2]], [[0]], [[0, 1], [0, 1, 2, 3, 4]]]
>>> structured_enumerate(x)
[[[0, 1, 2], [3, 4, 5]],
 [[6, 7], [8, 9, 10]],
 [[11]],
 [[12, 13], [14, 15, 16, 17, 18]]]

~edit~

Here's an attempt at a generic solution that works with any iterable, indexable or not, where you can specifiy iterable types to exclude from iteration.

from itertools import count

def structured_enumerate(iterable, dontiter=(str,), counter=None):
    'enumerate elements in nested iterable, preserve structure'
    result = []
    if counter is None:
        counter = count()

    for x in iterable:
        # check if x should be iterated
        try:
            iter(x)
            is_iterable = True
        except TypeError:
            is_iterable = False

        # strings of length zero and one are a special case
        if isinstance(x, str) and len(x) < 2:
            is_iterable = False

        if is_iterable and not isinstance(x, dontiter):
            subresult = structured_enumerate(x, dontiter, counter)
            result.append(subresult)
        else:
            result.append(next(counter))

    return result

Demo:

>>> fuzzy = [{0, 0}, '000', [0, [0, 0]], (0,0), 0]
>>> structured_enumerate(fuzzy)
[[0, 1], 2, [3, [4, 5]], [6, 7], 8]
>>> structured_enumerate(fuzzy, dontiter=())
[[0, 1], [2, 3, 4], [5, [6, 7]], [8, 9], 10]
>>> structured_enumerate(fuzzy, dontiter=(tuple, set))
[0, [1, 2, 3], [4, [5, 6]], 7, 8]