Search code examples
pythonlistlist-comprehension

How to return a subset of a list that matches a condition


Let's say I have a list of ints:

listOfNumbers = range(100)

And I want to return a list of the elements that meet a certain condition, say:

def meetsCondition(element):
    return bool(element != 0 and element % 7 == 0)

What's a Pythonic way to return a sub-list of element in a list for which meetsCondition(element) is True?

A naive approach:

def subList(inputList):
    outputList = []

    for element in inputList:
        if meetsCondition(element):
            outputList.append(element)

    return outputList

divisibleBySeven = subList(listOfNumbers)

Is there a simple way to do this, perhaps with a list comprehension or set() functions, and without the temporary outputList?


Solution

  • Use list comprehension,

    divisibleBySeven = [num for num in inputList if num != 0 and num % 7 == 0]
    

    or you can use the meetsCondition also,

    divisibleBySeven = [num for num in inputList if meetsCondition(num)]
    

    you can actually write the same condition with Python's truthy semantics, like this

    divisibleBySeven = [num for num in inputList if num and num % 7]
    

    alternatively, you can use filter function with your meetsCondition, like this

    divisibleBySeven = filter(meetsCondition, inputList)
    

    %timeit

    listOfNumbers = range(1000000)
    
    %timeit [num for num in listOfNumbers if meetsCondition(num)]
    [out]:
    243 ms ± 4.51 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
    
    %timeit list(filter(meetsCondition, listOfNumbers))
    [out]:
    211 ms ± 4.19 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)