CODEWITHSUNDEEP
pythonpython-3.xlist-comprehension
Filip
Filip

Reputation: 806

Python: List comprehension with list of lists

How would I write the following code in list comprehension?

with open('some_file.txt', 'r') as f:
    lines = f.readlines()

lines = [line.strip('\n') for line in lines]
list_of_lists = [[int(elm) for elm in line.split(' ')] for line in lines]

My file looks something like this:

3 8 6 9 4
4 3 0 8 6
2 8 3 6 9
3 7 9 0 3

Which means that:

grid = '3 8 6 9 4\n4 3 0 8 6\n2 8 3 6 9\n3 7 9 0 3'

Upvotes: 2

Views: 1579

Answers (4)

jpp
jpp

Reputation: 164643

You don't need to apply str.strip and str.split separately. Instead, combine them in one operation. The list comprehension is constructed by defining a list element followed by an iteration over a for loop.

Also note that str.strip without an argument will deal with \n as well as whitespace. Similarly, str.split without an argument will split by whitespace.

from io import StringIO

x = StringIO("""3 8 6 9 4
4 3 0 8 6
2 8 3 6 9
3 7 9 0 3""")

# replace x with open('some_file.txt', 'r')
with x as grid:
    list_of_lists = [[int(elm) for elm in line.strip().split()] for line in grid]

Result:

print(list_of_lists)

[[3, 8, 6, 9, 4],
 [4, 3, 0, 8, 6],
 [2, 8, 3, 6, 9],
 [3, 7, 9, 0, 3]]

With built-ins, it's marginally more efficient to use map:

list_of_lists = [list(map(int, line.strip().split())) for line in grid]

Upvotes: 1

Rafał
Rafał

Reputation: 705

You can use map function

grid = open('some.txt', 'r')
lines = [line.strip('\n') for line in grid]

list_of_lists = []

f = lambda line: [int(elm) for elm in line.split(' ')]
list_of_lists = map(f, lines)
print(list_of_lists)

Or You could implement it even shorter

f = lambda line: map(int, line.split(' '))
list_of_lists = map(f, lines)
print(list_of_lists)

Upvotes: 0

Vikas Periyadath
Vikas Periyadath

Reputation: 3186

Here a try , first split each lines and you will get a list of numbers as string, so map function can be used to change it to int:

with open('file.txt', 'r') as f:
    k = [list(map(int,i.split())) for i in f.readlines()]
    print(k)

Upvotes: 2

timgeb
timgeb

Reputation: 78680

int does not care about newline characters, so you can skip stripping those.

>>> fake_file = '3 8 6 9 4\n4 3 0 8 6\n2 8 3 6 9\n3 7 9 0 3\n'.splitlines()
>>> [[int(x) for x in line.split()] for line in fake_file]
[[3, 8, 6, 9, 4], [4, 3, 0, 8, 6], [2, 8, 3, 6, 9], [3, 7, 9, 0, 3]]

Note that pandas and numpy offer high level functions for reading in data such as this, e.g. numpy.loadtxt or pandas.read_csv.

Upvotes: 1

Related Questions