list comprehension replace for loop in 2D matrix

Question

I try to use list comprehension to replace the for loop.

original file is

2 3 4 5 6 3
1 2 2 4 5 5
1 2 2 2 2 4

for loop

line_number = 0
for line in file:
    line_data = line.split()
    Cordi[line_number, :5] = line_data 
    line_number += 1

output is

[[2 3 4 5 6 3]
 [1 2 2 4 5 5]
 [1 2 2 2 2 4]]

if use list comprehension instead, for what I can think of is (I have to change the data type to int, so it can be plotted in later part of the program)

Cordi1= [int(x) for x in line.split() for line in data]

but the output is

[1, 1, 1]

but line.split() for line in data is actually a list, and if I try

Cordi1 = [int(x) for x in name of the list]

it works, why this happens?

You can produce a list with a list comprehension, but `Cordi` is a numpy array instead. — Martijn Pieters, Aug 17 '14 at 02:34
@MartijnPieters, ah, so I cannot produce a 2D array with list comprehension? — rankthefirst, Aug 17 '14 at 13:29
You could convert the nested lists to a 2D array afterwards, but there could well be better options for producing a numpy array from this data. — Martijn Pieters, Aug 17 '14 at 15:15
See [how to convert 2d list to 2d numpy array?](http://stackoverflow.com/q/7717380) for a conversion option. — Martijn Pieters, Aug 17 '14 at 15:19
@MartijnPieters Thank you, ah, but actually either a nested 2D list or a numpy 2D list, they work well plotted. But in future applications, that may cause problems. Thank you anyway! — rankthefirst, Aug 17 '14 at 18:45

Martijn Pieters · Accepted Answer · 2014-08-17T16:11:23.277

You have the order of your loops swapped; they should be ordered in the same way they would be nested, from left to right:

[int(x) for line in data for x in line.split()]

This loops over data first, then for each line iteration, iterates over line.split() to produce x. You then produce one flat list of integers from these.

However, since you are trying to build a list of lists, you need to nest a list comprehension inside another:

Cordi1 = [[int(i) for i in line.split()] for line in data]

Demo:

>>> data = '''\
... 2 3 4 5 6 3
... 1 2 2 4 5 5
... 1 2 2 2 2 4
... '''.splitlines()
>>> [int(x) for line in data for x in line.split()]
[2, 3, 4, 5, 6, 3, 1, 2, 2, 4, 5, 5, 1, 2, 2, 2, 2, 4]
>>> [[int(i) for i in line.split()] for line in data]
[[2, 3, 4, 5, 6, 3], [1, 2, 2, 4, 5, 5], [1, 2, 2, 2, 2, 4]]

If you wanted a multidimensional numpy array from this, you can either convert the above directly to an array or create an array from the data then reshape:

>>> import numpy as np
>>> np.array([[int(i) for i in line.split()] for line in data])
array([[2, 3, 4, 5, 6, 3],
       [1, 2, 2, 4, 5, 5],
       [1, 2, 2, 2, 2, 4]])
>>> np.array([int(i) for line in data for i in line.split()]).reshape((3, 6))
array([[2, 3, 4, 5, 6, 3],
       [1, 2, 2, 4, 5, 5],
       [1, 2, 2, 2, 2, 4]])

list comprehension replace for loop in 2D matrix

1 Answers1

Linked