CODEWITHSUNDEEP
pythonfor-loop
Eial Teomy
Eial Teomy

Reputation: 11

How do I rewrite the following code in Python without for loops?

I have an array "b" of size L^2 x L*(L+1), and an array "a" of size L x L.

currently my code is

for i in range (L):
    for j in range (L):
        b[i+j*L,i+j*(L+1)] = a[j,i]

What this means is that, for example for L=2, if the 2x2 array "a" has the form

ab
cd

then I want the 4x6 array "b" to be

a00000
0b0000
000c00
0000d0

How do I rewrite the same thing without using for loops?

Upvotes: 1

Views: 175

Answers (2)

Milo
Milo

Reputation: 2171

What you want is to fill the diagonal of matrix B with the values of (flattened) A. Numpy has functions for this:

https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flatten.html

https://numpy.org/doc/stable/reference/generated/numpy.fill_diagonal.html

import numpy as np

# Set sample data
a = np.array([[1, 2], [3, 4]])
b = np.zeros([4,6])

# This is it:
np.fill_diagonal(b, a.flatten())

If you don't want to use a library, for example because this is a programming assignment, you can represent matrices as nested lists and use a list comprehension, as this:

# Prepare data
a = [[1, 2], [3, 4]]
L = len(a)

# "Build" the result
b = [[a[i//L][i%L] if i == j else 0 for j in range(L*(L+1))] for i in range(L*L)]

# Same, with better formatting:
b = [[a[i//L][i%L] 
      if i == j else 0
      for j in range(L*(L+1))]
     for i in range(L*L)]

# b will be = [[1, 0, 0, 0, 0, 0], 
#              [0, 2, 0, 0, 0, 0], 
#              [0, 0, 3, 0, 0, 0], 
#              [0, 0, 0, 4, 0, 0]]

Anyway you need to iterate through the items in 'a', so you are just replacing the 'for' constructions by a list comprehension. This might be more efficient for large matrices but arguably less clear.

Upvotes: 2

Serial Lazer
Serial Lazer

Reputation: 1669

Generalizing the answer from Milo:

L = (a.shape)[0]
b = np.zeros([L*L, L*(L+1)])

np.fill_diagonal(b, a.flatten())

Upvotes: 1

Related Questions