Search code examples
pythoniterationlist-comprehension

How can I collect the results of a repeated calculation in a list, dictionary etc. (or make a copy of a list with each element modified)?


There are a great many existing Q&A on Stack Overflow on this general theme, but they are all either poor quality (typically, implied from a beginner's debugging problem) or miss the mark in some other way (generally by being insufficiently general). There are at least two extremely common ways to get the naive code wrong, and beginners would benefit more from a canonical about looping than from having their questions closed as typos or a canonical about what printing entails. So this is my attempt to put all the related information in the same place.

Suppose I have some simple code that does a calculation with a value x and assigns it to y:

y = x + 1

# Or it could be in a function:
def calc_y(an_x):
    return an_x + 1

Now I want to repeat the calculation for many possible values of x. I know that I can use a for loop if I already have a list (or other sequence) of values to use:

xs = [1, 3, 5]
for x in xs:
    y = x + 1

Or I can use a while loop if there is some other logic to calculate the sequence of x values:

def next_collatz(value):
    if value % 2 == 0:
        return value // 2
    else:
        return 3 * value + 1

def collatz_from_19():
    x = 19
    while x != 1:
        x = next_collatz(x)

The question is: how can I collect these values and use them after the loop? I tried printing the value inside the loop, but it doesn't give me anything useful:

xs = [1, 3, 5]
for x in xs:
    print(x + 1)

The results show up on the screen, but I can't find any way to use them in the next part of the code. So I think I should try to store the values in a container, like a list or a dictionary. But when I try that:

xs = [1, 3, 5]
for x in xs:
    ys = []
    y = x + 1
    ys.append(y)

or

xs = [1, 3, 5]
for x in xs:
    ys = {}
    y = x + 1
    ys[x] = y

After either of these attempts, ys only contains the last result.


Solution

  • General approaches

    There are three ordinary ways to approach the problem: by explicitly using a loop (normally a for loop, but while loops are also possible); by using a list comprehension (or dict comprehension, set comprehension, or generator expression as appropriate to the specific need in context); or by using the built-in map (results of which can be used to construct a list, set or dict explicitly).

    Using an explicit loop

    Create a list or dictionary before the loop, and add each value as it's calculated:

    def make_list_with_inline_code_and_for():
        ys = []
        for x in [1, 3, 5]:
            ys.append(x + 1)
        return ys
    
    def next_collatz(value):
        if value % 2 == 0:
            return value // 2
        else:
            return 3 * value + 1
    
    def make_dict_with_function_and_while():
        x = 19
        ys = {}
        while x != 1:
            y = next_collatz(x)
            ys[x] = y # associate each key with the next number in the Collatz sequence.
            x = y # continue calculating the sequence.
        return ys
    

    In both examples here, the loop was put into a function in order to label the code and make it reusable. These examples return the ys value so that the calling code can use the result. But of course, the computed ys could also be used later in the same function, and loops like these could also be written outside of any function.

    Use a for loop when there is an existing input, where each element should be processed independently. Use a while loop to create output elements until some condition is met. Python does not directly support running a loop a specific number of times (calculated in advance); the usual idiom is to make a dummy range of the appropriate length and use a for loop with that.

    Using a comprehension or generator expression

    A list comprehension gives elegant syntax for creating a list from an existing sequence of values. It should be preferred where possible, because it means that the code does not have to focus on the details of how to build the list, making it easier to read. It can also be faster, although this will usually not matter.

    It can work with either a function call or other calculation (any expression in terms of the "source" elements), and it looks like:

    xs = [1, 3, 5]
    
    ys = [x + 1 for x in xs]
    # or
    def calc_y(an_x):
        return an_x + 1
    ys = [calc_y(x) for x in xs]
    

    Note that this will not replace a while loop; there is no valid syntax replacing for with while here. In general, list comprehensions are meant for taking existing values and doing a separate calculation on each - not for any kind of logic that involves "remembering" anything from one iteration to the next (although this can be worked around, especially in Python 3.8 and later).

    Similarly, a dictionary result can be created using a dict comprehension - as long as both a key and value are computed in each iteration. Depending on exact needs, set comprehensions (produce a set, which does not contain duplicate values) and generator expressions (produce a lazily-evaluated result; see below about map and generator expressions) may also be appropriate.

    Using map

    This is similar to a list comprehension, but even more specific. map is a built-in function that can apply a function repeatedly to multiple different arguments from some input sequence (or multiple sequences).

    Getting results equivalent to the previous code looks like:

    xs = [1, 3, 5]
    
    def calc_y(an_x):
        return an_x + 1
    
    ys = list(map(calc_y, xs))
    # or
    ys = list(map(lambda x: x + 1, xs))
    

    As well as requiring an input sequence (it doesn't replace a while loop), the calculation needs to be done using a function or other callable, such as the lambda shown above (any of these, when passed to map, is a so-called "higher-order function").

    In Python 3.x, map is a class, and calling it therefore creates an instance of that class - and that instance is a special kind of iterator (not a list) that can't be iterated more than once. (We can get something similar using a generator expression rather than a list comprehension; simply use () instead of [].)

    Therefore, the code above explicitly creates a list from the mapped values. In other situations, it might not be necessary to do this (i.e., if it will only be iterated over once). On the other hand, if a set is necessary, the map object can be passed directly to set rather than list in the same way. To produce a dictionary, the map should be set up so that each output element is a (key, value) tuple; then it can be passed to dict, like so:

    def dict_from_map_example(letters):
        return dict(map(lambda l: (l, l.upper()), letters))
        # equivalent using a dict comprehension:
        # return {l:l.upper() for l in letters}
    

    Generally, map is limited and uncommon compared to list comprehensions, and list comprehensions should be preferred in most code. However, it does offer some advantages. In particular, it can avoid the need to specify and use an iteration variable: when we write list(map(calc_y, xs)), we don't need to make up an x to name the elements of xs, and we don't have to write code to pass it to calc_y (as in the list comprehension equivalent, [calc_y(x) for x in xs] - note the two xs). Some people find this more elegant.