Home > Net >  Turn python code into a generator function
Turn python code into a generator function

Time:12-10

How can I turn this code into a generator function? Or can I do it in a other way avoiding reading all data into memory? The problem right now is that my memory gets full. I get KILLED after a long time when executing the code.

Code:

data = [3,4,3,1,2]

def convert(data):
    for index in range(len(data)):
        if data[index] == 0:
            data[index] = 6
            data.append(8)
        elif data[index] == 1:
            data[index] = 0
        elif data[index] == 2:
            data[index] = 1
        elif data[index] == 3:
            data[index] = 2
        elif data[index] == 4:
            data[index] = 3
        elif data[index] == 5:
            data[index] = 4
        elif data[index] == 6:
            data[index] = 5
        elif data[index] == 7:
            data[index] = 6
        elif data[index] == 8:
            data[index] = 7

    return data

for i in range(256):
    output = convert(data)
    print(len(output))

Output:

266396864
290566743
316430103
346477329
376199930
412595447
447983143
490587171
534155549
582826967
637044072
692630033
759072776
824183073
903182618
982138692
1073414138
1171199621
1275457000
1396116848
1516813106
Killed

CodePudding user response:

To answer the question: to turn a function into a generator function, all you have to do is yield something. You might do it like this:

def convert(data):
    for index in range(len(data)):
        ...

        yield data

Then, you can iterate over the output like this:

iter_converted_datas = convert(data)

for _, converted in zip(range(256), iter_converted_datas):
    print(len(converted))

However, I would suggest some improvements to this code.

First of all, if you intend to iterate over the entire list length, use enumerate:

for _, converted in enumerate(iter_converted_datas):

The next thing to do is get rid of all those elif statements.

One helpful thing might be to supply a dictionary argument to your generator function that tells it how to convert the data values (the first one is a special case since it also appends).

Here is what that dict might look like:

replacement_dict = {
    0: 6,
    1: 0,
    2: 1,
    3: 2,
    4: 3,
    5: 4,
    6: 5,
    7: 6,
    8: 7,
}

By the way: replacing a series of elif statements with a dictionary is a pretty typical thing to do in python. It isn't always appropriate, but it often works well.

Now you can write your generator like this:

def convert(data, replacement_dict):
    for index in range(len(data)):
        if index==0:
            lst.append(8)
        data[index] = replacement_dict[index]

And use it like this:

iter_converted_datas = convert(data, replacement_dict)

for _, converted in enumerate(iter_converted_datas):
    print(len(converted))

Now, it's possible you might have several indexes that you want to do other modifications with. You might split these up into multiple dictionaries with different purposes: a replacement_dict and a modification_dict.

The modification dict looks like this (it only has on entry right now, for index 0):

modification_dict = {
    0: lambda data: data.append(8)
}

Now, we will modify the generator function to accept both a replacement_dict and a modification_dict:

def convert(data, replacement_dict, modification_dict):
    for index in range(len(data)):
        # this modifies the data
        try:
            modification_dict[index](data)
        except KeyError:
            pass
        # this replaces the data
        data[index] = replacement_dict[index]

Now, use it like so:

iter_converted_datas = convert(data, replacement_dict, modification_dict)

for _, converted in enumerate(iter_converted_datas):
    print(len(converted))

CodePudding user response:

Here are a few things you can do to optimize it:

Instead of range(len(data)) you can use enumerate(data). This gives you access to both the element AND it's index. Example:

EDIT: According to this post, range is faster than enumerate. If you care about speed, you could ignore this change.

for index, element in enumerate(data):
    if element == 0:
        data[index] = 6

Secondly, most of the if statements have a predictable pattern. So you can rewrite them like this:

def convert(data):
    for idx, elem in enumerate(data):
        if elem == 0:
            data[idx] = 6
            data.append(8)

        if elem <= 8:
            data[index] = elem - 1

Since lists are mutable, you don't need to return data. It modifies it in-place.

  • Related