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

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]

CodePudding user response：

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)

CodePudding user response：

You can use str.join(iterable) inside a f-string to remove the single quote characters around the string elements when you print:

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]

CodePudding user response：

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.