user258532
user258532

Reputation: 541

Pass several arguments to function from map()

For example, if I want to detect all odd numbers in an array and set them to zero, I can use:

def setToZeroIfOdd(n):
   if n % 2 == 0:
      pass
   else:
      return 0

numbers = range(1,1000)

numbers = map(setToZeroIfOdd, numbers)

which works like a charm.

But when I try something like

def setToZeroIfDivisibleBy(n, divisor):
  if n % divisor == 0:
     return 0
  else:
     pass

numbers = map(setToZeroIfDivisibleBy(divisor=3), numbers)

it expects two arguments. Likewise,

numbers = map(setToZeroIfDivisibleBy, numbers, divisor=3)

does not work. How can I pass that divisor argument from within map()?

Upvotes: 6

Views: 6717

Answers (4)

chepner
chepner

Reputation: 531075

Another approach, instead of using partial, is to supply an infinite (or at least, long enough) sequence of 2nd arguments for the two-argument function:

from itertools import repeat
numbers = map(setToZeroIfDivisibleBy, numbers, repeat(3))

In Python 2, map will append None as necessary to the shorter of the two sequences to make them the same length. Assuming that will cause problems (either because your function cannot handle None as an input value or you end up with an infinite loop), you can either use itertools.imap, which stops after exhausting the shorter sequence:

from itertools import imap, repeat
numbers = list(imap(setToZeroIfDivisibleBy, numbers, repeat(3)))

or pass the length of numbers as a second argument to repeat so that the two sequences are the same length.

from itertools import repeat
numbers = map(setToZeroIfDivisibleBy, numbers, repeat(3, len(numbers)))

Upvotes: 0

Patrick Haugh
Patrick Haugh

Reputation: 60974

You can use functools.partial to make partial functions

from functools import partial

def setToZeroIfDivisibleBy(n, divisor):
  if n % divisor == 0:
     return 0
  else:
     pass

numbers = range(1,1000)

numbers = map(partial(setToZeroIfDivisibleBy, divisor=3), numbers)

Upvotes: 11

deceze
deceze

Reputation: 522042

You make a function which returns a function:

def setToZeroIfDivisibleBy(divisor):
    def callback(n):
        if n % divisor == 0:
            return 0
        else:
            pass

    return callback

numbers = map(setToZeroIfDivisibleBy(3), numbers)

BTW, you can entirely omit empty branches like else: pass; it doesn't do anything. Since it results in a None, I don't think that's what you want either. You probably want return n there instead.

Upvotes: 8

damisan
damisan

Reputation: 1047

Try using lambda function

numbers = map(lambda n: setToZeroIfDivisibleBy(n, divisor=3), numbers)

And rather than pass did you mean return n?

Upvotes: 10

Related Questions