Convert carry-add loop to a map or reduce to one-liner

Question

I want to perform standard additive carry on a vector. The base is a power of 2, so we can swap modulus for bitwise AND.

def carry(z,direction='left',word_size=8):
 v = z[:]
 mod = (1<<word_size) - 1
 if direction == 'left': v = v[::-1]
 accum = 0
 for i in xrange(len(v)):
 v[i] += accum
 accum = v[i] >> word_size
 v[i] = v[i] & mod
 print accum,v
 if direction == 'left': v = v[::-1]
 return accum,v

Is there any way to make this function even tinier?

• tokland: It would be very handy if you include some asserts (3 ó 4) so people can test their refactors easily.
• Cris Stringfellow: Good point. I will.

Answer 1

1. Your code has a lot of copying in it. In the default case (where direction is left) you copy the array three times. This seems like a lot. There are various minor improvements you can make, for example instead of

   v = v[::-1]
   

   you can write

   v.reverse()
   

   which at least re-uses the space in v.

   But I think that it would be much better to reorganize the whole program so that you store your bignums the other way round (with the least significant word at the start of the list), so that you can always process them in the convenient direction.

2. The parameters direction and word_size are part of the description of the data in z. So it would make sense to implement this code as a class to keep these values together. And then you could ensure that the digits always go in the convenient direction for your canonicalization algorithm.

   class Bignum(object):
    def __init__(self, word_size):
    self._word_size = word_size
    self._mod = 2 ** word_size
    self._digits = []
    def canonicalize(self, carry = 0):
    """
    Add `carry` and canonicalize the array of digits so that each
    is less than the modulus.
    """
    assert(carry >= 0)
    for i, d in enumerate(self._digits):
    carry, self._digits[i] = divmod(d + carry, self._mod)
    while carry:
    carry, d = divmod(carry, self._mod)
    self._digits.append(d)
   

3. Python already has built-in support for bignums, anyway, so what exactly are you trying to do here that can't be done in vanilla Python?

   Edited to add: I see from your comment that I misunderstood the context in which this function would be used (so always give us the context!). It's still the case that you could implement what you want using Python's built-in bignums, for example if you represented your key as an integer then you could write something like:

   def fold(key, k, word_size):
    mod = 2 ** (k * word_size)
    accum = 0
    while key:
    key, rem = divmod(key, mod)
    accum += rem
    return accum % mod
   

   but if you prefer to represent the key as a list of words, then you could still implement the key-folding operation directly:

   from itertools import izip_longest
   class WordArray(object):
    def __init__(self, data = [], word_size = 8):
    self._word_size = word_size
    self._mod = 1 << word_size
    self._data = data
    def fold(self, k):
    """
    Fold array into parts of length k, add them with carry, and return a
    new WordArray. Discard the final carry, if any.
    """
    def folded():
    parts = [xrange(i * k, min((i + 1) * k, len(self._data))) 
    for i in xrange((len(self._data) + 1) // k)]
    carry = 0
    for ii in izip_longest(*parts, fillvalue = 0):
    carry, z = divmod(sum(self._data[i] for i in ii) + carry, self._mod)
    yield z
    return WordArray(list(folded()), self._word_size)
   

• Cris Stringfellow: Thanks a lot. I will make my way through it. Just quickly the objective is not bignum. I fold an array into parts of length k and add these together, but I preserve the length k, so I discard any carry that falls off one end. It's a mixing step in a key scheduling part of a crypto.
• Jaime: Your approach reminded me of this: pyvideo.org/video/880/stop-writing-classes 'the signature of this shouldn't be a class is that it has two methods, one of which is __init__, anytime you see this you should think "hey, maybe I only need that one method!"'
• Gareth Rees: @Jaime: it should be clear, I hope, that my classes aren't intended to be complete. I believe that the OP has more operations that he hasn't told us about, that in a complete implementation would become methods on these classes. (Do you think I need to explain this point?)
• Cris Stringfellow: Had a closer, look @GarethRees -- that's awesome. I really like that. Very concise on the loops in the BigNum class. Looking at the next part now.
• Cris Stringfellow: You understood me perfectly on the second part. I am pretty shocked since I gave just a tiny description. That is really clever, too. You are amazing.

Answer 2

I got it:

def carry2(z,direction='left',word_size=8):
 x = z[:]
 v = []
 mod = (1<<word_size) - 1
 if direction == 'left': x.reverse()
 def cu(a,n):
 v.append((a+n)&mod)
 return (a+n) >> word_size
 accum = reduce(cu,x,0)
 if direction == 'left': v.reverse()
 return accum,v

• tokland: You are clearly stuck in the imperative paradigm, that's why the code is so verbose (and hard to understand). Check docs.python.org/2/howto/functional.html and google for "python functional programming".
• Cris Stringfellow: Thanks! I'm getting started. My latest code has a lot more generators, lambdas and functional style...but I am just getting started.

Answer 3

Preliminary remarks :

• Functional idioms and imperative mutations doesn't play well (except when they do. Cf Scala). Here, both FP and Python lovers will likely be lost !
• Since there is exactly 2 possible directions, use a Boolean. It will prevent typo (trying to pass 'Left' for instance).
• Extract the main computation to a dedicated function. First benefice : you won't have to test for direction twice, which is error prone.
• (1st version) : instead of copying + updating via index (awkward), generate a brand new sequence.

That's an interesting question, since it requires both folding (to propagate the carry) and mapping from updated values (to apply mod, ie get the less significant bits). Since it's not that common, I suggest to stick with a standard Python approach. Don't force an idiom if it obfuscates.

def carry_from_left(z, word_size):
 mod = (1<<word_size) - 1
 res = []
 acc = 0
 for a in z :
 tot = a + acc
 res.append(tot & mod)
 acc = tot >> word_size
 return (acc, res)
def carry(z, left_to_right = False, word_size=8):
 if left_to_right :
 return carry_from_left(z, word_size)
 else:
 acc, res = carry_from_left(reversed(z), word_size)
 res.reverse() #inplace. Nevermind, since it's a brand new list.
 return (acc, res)

Now, if this "map+fold" pattern occurs frequently, you can abstract it.

def fold_n_map (f, l, acc):
 """ Apply a function (x, acc) -> (y, acc') cumulatively.
 Return a tuple consisting of folded (reduced) acc and list of ys.
 TODO : handle empty sequence, optional initial value, etc,
 in order to mimic 'reduce' interface"""
 res = []
 for x in l :
 y, acc = f(x, acc)
 res.append(y)
 return acc, res
def carry_fold_left(z, word_size):
 mod = (1<<word_size) - 1
 #Nearly a one-liner ! Replace lambda by named function for your Python fellows.
 return fold_n_map(lambda x, acc: ((x+acc) & mod, (x+acc) >> word_size), z, 0)
def carry(z, left_to_right = False, word_size=8):
 #idem
 ...

• Cris Stringfellow: Cool, I like that. Especially the abstraction of the pattern.