Is there a pythonic way to re-write the following block of code? List comprehensions?

Question

I am trying to do this . This is what I came up with:

X, Y, A = [x1], [y1], [(x1 + y1) % F]
for i in range(1, N):
  X.append((C * X[i-1] + D * Y[i-1] + E1) % F)
  Y.append((D * X[i-1] + C * Y[i-1] + E2) % F)
  A.append((X[i] + Y[i]) % F)

This work but as you see it's not very pythonic. I was wondering if there is a way to get the same result in a pythonic way, probably by using list comprehension(s). Thanks!

Answer 1

If you have two variables, each of which should be updated using the current value of both variables, the neat solution is to use a parallel assignment.

X, Y, A = [], [], []
x, y = x1, y1
for i in range(N):
    X.append(x)
    Y.append(y)
    A.append((x + y) % F)
    # update both x and y in parallel
    x, y = (C*x + D*y + E1) % F, (D*x + C*y + E2) % F

Answer 2

Given x , and y don't depend on A , I might be tempted to create a generator for them, and then build X , Y , and A from that (or do whatever).

def genXYA(n):
    x, y = x1, y1
    for i in range(n):
        nx, ny = (C * x + D * y + E1) % F, (D * x + C * y + E2) % F
        a = (nx + ny) % F
        yield nx, ny, a
        x, y = nx, ny


X, Y, A = map(list, zip(*genXYA(N)))

---

Just to explain what's happening here...

A generator is like a list, except rather than returning everything immediately, it only returns one item at a time.

The function genXYA(n) returns a generator which returns tuples (x, y, a) at every step, and stops after n things have been returned.

Now we have this generator, we can directly create the lists X , Y , and A .

Say N = 5 , if you did list(genXYA(N)) , you'd get

[(x2, y2, a1),
 (x3, y3, a2),
 (x4, y4, a3),
 (x5, y5, a4),
 (x6, y6, a5)]

And what you want is to separate those columns into individual named lists. You can use zip for this. zip takes a number of iterables, and creates new iterables which are the first of each, the second of each, the third of each, etc.

In order to give the right argument to zip we need to unpack the generator. That's what the * is for.

The result of zip(*genXYA(N)) is another generator, but we can see what's inside by calling list(zip(*genXYA(N))) .

[(x2, x3, x4, x5, x6), # => X
 (y2, y3, y4, y5, y6), # => Y
 (a1, a2, a3, a4, a5)] # => A

So we're getting closer the end result.

When we call map(list, zip(*genXYA(N))) , we're saying that for everything in that list, we want a list of each item (and a list of a tuple is a list containing everything in the tuple).

map also returns a generator, which we can then unpack into the variables we want.

---

The benefit of using generators is that you only calculate the next iteration when you need it, and you don't have to store all previous iterations.

For example, if you just wanted the last value of A , you could do:

for x, y, a in genXYA(N):
    pass
print(a)

Answer 3

Using recursion+ numpy :

import numpy as np

X=np.array([[x1, y1, 1]])
A=np.array([[C, D, E1], [D, C, E2], [1,1,1]])

def myFunc(res, A, F, N):
    res[-1,2]=1
    if(N==1):
        return np.append(res, [A@res[-1]%F], axis=0)
    else:
        N-=1
        return myFunc(np.append(res, [A@res[-1]%F], axis=0), A, F, N)


X, Y=myFunc(X, A, F, N-1)[:, :2].T

A=(X+Y)%F

Sample output:

x1=1
y1=2
C, D, E1, E2, F, N=4,15,6,7,3,11

X=np.array([[x1, y1, 1]])
A=np.array([[C, D, E1], [D, C, E2], [1,1,1]])

X, Y=myFunc(X, A, F, N-1)[:, :2].T
A=(X+Y)%F

print(X)
print(Y)
print(A)
 #outputs:

[1 1 1 1 1 1 1 1 1 1 1]
[2 0 1 2 0 1 2 0 1 2 0]
[0 1 2 0 1 2 0 1 2 0 1]

Answer 4

You can use itertools.accumulate to express this recurrence relation and generate both x's and y's in a tuple.

import itertools
# don't actually care about entries outside the first, but needs to be right length I believe
initial_array = [(x1, y1)] + list(itertools.repeat((0,0), N-1))
XY = itertools.accumulate(initial_array, lambda prev, _: ((C*prev[0]+D*prev[1]+E1)%F, (D*prev[0]+C*prev[1]+E2)%F))
A = [(x+y)% F for x, y in XY]

This should work. XY is a generator, so you also gain the benefits of not needing to store the full list in memory.

Answer 5

This is what I ended up writing. Since I don't need to store x and y values:

def generate_xy(x, y, n):
  yield x, y
  for _ in range(2, n+1):
    next_x = (C*x + D*y + E1) % F
    next_y = (D*x + C*y + E2) % F
    yield next_x, next_y
    x, y = next_x, next_y

A = [(x+y)%F for x, y in generate_xy(x1, y1, N)]

Python is quite the language!