Python - Check if element is in 1 or more nested lists and return whole list if TRUE

Question

Thanks in advance for any help.

I have a nested list/list of lists and I need to return each list to the screen ONLY if the list contains a specific element.

The list characteristic:

• Each list has the same amount of elements,
• Each element represents the same thing
• In the below example index 0 = fruit name, index 1 = amount, index 2 = colour .
• Example: lists = [banana, 10, yellow], [apple, 12, red], [pear, 60, green], [mango, 5, yellow]

I have played around with conditional for loops and with the idea of just creating a new list for each variation but that seems like a hard to manage solution.

Can anyone help?

Search 1: If 'banana' is in one or more nested lists Then print each nested list

Expected output: [banana, 10, yellow]

Search 2: If 'yellow' is in one or more nested lists Then print each nested list

Expected output: [banana, 10, yellow] [mango, 5, yellow]

Answer 1

Here is a rough way to do it:

lists = [
    ["banana", 10, "yellow"], 
    ["apple", 12, "red"], 
    ["pear", 60, "green"], 
    ["mango", 5, "yellow"],
]

keyword = 'banana'
for lst in lists:
    if keyword in lst:
        print(lst)

keyword = 'yellow'
for lst in lists:
    if keyword in lst:
        print(lst)

Ideally you would extract the search to a function accepting the lists and the keyword:

def get_sublists_containing_keyword(lists, keyword):
    sublists = []
    for lst in lists:
        if keyword in lst:
            sublists.append(lst)
    return sublists

lists = [
    ["banana", 10, "yellow"], 
    ["apple", 12, "red"], 
    ["pear", 60, "green"], 
    ["mango", 5, "yellow"],
]

banana_lists = get_sublists_containing_keyword(lists, 'banana')
yellow_lists = get_sublists_containing_keyword(lists, 'yellow')

for banana_list in banana_lists:
    print(banana_list)
for yellow_list in yellow_lists:
    print(yellow_list)

Answer 2

You can use str. inside a f-string to remove the single quote characters around the string elements when you print: inside a f-string以在打印时删除字符串元素周围的单引号字符：

def print_lists_that_contain_search_term(lists: list[list[str | int]], 
                                         search_term: str) -> None:
    print(f'{search_term = }')
    print(' '.join(f'[{", ".join(map(str, lst))}]' for lst in lists if search_term in lst))

def main() -> None:
    lists = [['banana', 10, 'yellow'], ['apple', 12, 'red'], ['pear', 60, 'green'], ['mango', 5, 'yellow']]
    print_lists_that_contain_search_term(lists, 'banana')
    print_lists_that_contain_search_term(lists, 'yellow')

if __name__ == '__main__':
    main()

Output:

search_term = 'banana'
[banana, 10, yellow]
search_term = 'yellow'
[banana, 10, yellow] [mango, 5, yellow]

Answer 3

Here is a single liner solution for this.

def search(lists, item):
    return list(filter(None, map(lambda x: x if item in x else [], lists)))

Now you can call the function and check it.

In [12]: lists
Out[12]: 
[['banana', 10, 'yellow'],
 ['apple', 12, 'red'],
 ['pear', 60, 'green'],
 ['mango', 5, 'yellow']]

In [13]: search(lists, 'banana')
Out[13]: [['banana', 10, 'yellow']]

In [14]: search(lists, 'yellow')
Out[14]: [['banana', 10, 'yellow'], ['mango', 5, 'yellow']]

Code explanation

Here I Have used lambda expression and checked if the tobe search item is inside the list then return that list else return an empty list. And removed all the empty lists via a filter Function.