CODEWITHSUNDEEP
python
João Palma
João Palma

Reputation: 159

Sum OR average OR max elements of two lists according to a third list defining the operation

I have 3 lists:

a = [1,2,3]
b = [4,5,6]
c = ['SUM','AVG','MAX']

I need to achieve:

d = [5,3.5,6]

Note that a and b are just an example. In reality I have a few hundred arrays with about one hundred elements. So I would prefer an interesting "Pythonian way" to achieve the d array, rather than looping around. Something like similar to map, for example:

d= [calculate(c, x) for x in [a,b]]

Is this possible?

Upvotes: 2

Views: 116

Answers (3)

Padraic Cunningham
Padraic Cunningham

Reputation: 180481

If you have less functions than elements in the lists you can cycle the functions:

a = [1, 2, 3]
b = [4, 5, 6]
c = ['SUM', 'AVG', 'MAX']

mapped = {'SUM': sum, 'AVG': lambda x: sum(x, 0.0) / len(x), 'MAX': max}

from itertools import cycle


def calc(funcs, *args):
    _fncs = list(map(mapped.get, funcs))
    return [f(it) for f, it in zip(cycle(_fncs), zip(*args))]

If c is shorter than a and b then the functions will be cycled, if the lists are different lengths you can use itertools.izip_longest but you will have to decide what an appropriate fillvalue is.

If you are using python 3 there is also a builtin for the average:

mapped = {'SUM': sum, 'AVG': mean, 'MAX': max}

from itertools import cycle

from statistics import mean

def calc(funcs, *args):
    _fncs = list(map(mapped.get, funcs))
    return [f(it) for f, it in zip(cycle(_fncs), zip(*args))]

If they are always going to be even lengths you don't need the cycle:

def calc(funcs, *args):
    _fncs = map(mapped.get, funcs)
    return [f(it) for f, it in zip(_fncs, zip(*args))]

Upvotes: 1

Tom Karzes
Tom Karzes

Reputation: 24062

Here's one way to do this:

op_dict = {
    'SUM' : lambda x, y: x + y,
    'AVG' : lambda x, y: (x + y) / 2.0,
    'MAX' : max
}

a = [1,2,3]
b = [4,5,6]
c = ['SUM','AVG','MAX']

d = [op_dict[z](x, y) for x, y, z in zip(a, b, c)]

Here's a more powerful version that can handle any number of lists, but they must still all be the same length:

op_dict = {
    'SUM' : lambda *x: sum(x),
    'AVG' : lambda *x: sum(x) / float(len(x)),
    'MAX' : max
}

ops = ['SUM', 'AVG', 'MAX']
a = [1, 2, 3]
b = [4, 5, 6]
c = [7, 8, 9]
d = [10, 11, 12]

args = [ops, a, b, c, d]

r = [op_dict[x[0]](*x[1:]) for x in zip(*args)]

Upvotes: 4

Deoxyribonucleic
Deoxyribonucleic

Reputation: 153

How about:

def sum(a, b):
    return a + b
def avg(a, b):
    return a + b / 2.0

def calculate(a, b, c):
    mapping = {'SUM': sum,
               'AVG': avg,
               'MAX': max}
    return mapping[c](a,b)

d = [calculate(a1,b1,c1) for a1,b1,c1 in zip(a, b, c)]

So, working from the bottom up:

  • Zip is great for iterating through multiple lists at once. a1,b1,c1 are the items in a, b, and c at the same index.
  • Calculate will do a different calculation depending on c.
  • the mapping dict simply stores functions ( notice the lack of parens, we don't call the function. We are storing the function itself. )
  • The we grab the correct function, determined by the key we pass in, and call it with two arguments.

edit: I love how we've all given essentially the same answer, within a small period of time. Kinda reassures you that this is the way to do it.

Obviously syntatic sugar like lambdas help keep the code looking simpler.

Upvotes: 0

Related Questions