在列表及其子列表中查找成员

Question

我想生成 function 以在列表中查找子列表：

data = ['expense', ['food', ['meal', 'snack', 'drink'], 'transportation', ['bus', 'railway']], 'income', ['salary', 'bonus']]
#for example if I want to find in 'food', I want the output to be [food, meal, snack , drink]

这是我到目前为止所拥有的：

def find_subcategories_gen(category, categories, found=False):
    if type(categories) == list:
        for index, child in enumerate(categories):
            yield from find_subcategories_gen(category, child, False)
            if child == category and index + 1 < len(categories) and type(categories[index + 1]) == list:
                # When the target category is found,
                # recursively call this generator on the subcategories
                # with the flag set as True.
                yield from find_subcategories_gen(category, categories[index+1], True)
    else:
        if categories == category or found == True:
            yield categories

这应该可以工作，但是当我尝试查找“食物”时，我得到的只是“食物”，而不是父列表和“食物”子列表。 我应该如何处理这个？

Answer 1

我添加了附加条件and found == False

data = ['expense', ['food', ['meal', 'snack', 'drink'], 'transportation', ['bus', 'railway']], 'income', ['salary', 'bonus']]

def find_subcategories_gen(category, categories, found=False):
    if type(categories) == list and found == False:  #added found==False
        for index, child in enumerate(categories):
            yield from find_subcategories_gen(category, child, False)
            if child == category and index + 1 < len(categories) and type(categories[index + 1]) == list:
                # When the target category is found,
                # recursively call this generator on the subcategories
                # with the flag set as True.
                yield from find_subcategories_gen(category, categories[index+1], True)
    else:
        if categories == category or found == True:
            yield categories
            
print(list(find_subcategories_gen('food', data)))

因此 output 是

['food', ['meal', 'snack', 'drink']]