For example, I have a list which only contains zeroes, ones and twos
ls = [
[1, 1, 0, 2, 2],
[1, 1, 1, 2, 2, 2],
[0, 1, 0, 0, 0, 2, 0]
]
I want to split this list it into two list,
ls1
contains the ones, and ls2
contains the twos. I would like to keep the shape and use 0 to replace 2's in ls1
and 1's in ls2
. The expected result is:
ls1 = [
[1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0]
]
ls2 = [
[0, 0, 0, 2, 2],
[0, 0, 0, 2, 2, 2],
[0, 0, 0, 0, 0, 2, 0]
]
I know I can use a for loop to handle it, but is there an elegant way?
ls1, ls2 = ([[b & x for b in a] for a in ls] for x in (1, 2))
Using nested list comprehension:
ls1 = [[1 if e == 1 else 0 for e in l] for l in ls]
ls2 = [[2 if e == 2 else 0 for e in l] for l in ls]
# ls1
[[1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0]]
# ls2
[[0, 0, 0, 2, 2],
[0, 0, 0, 2, 2, 2],
[0, 0, 0, 0, 0, 2, 0]]
Nothing especially "elegant" about it, but you could simply use a list comprehension:
ls = [
[1, 1, 0, 2, 2],
[1, 1, 1, 2, 2, 2],
[0, 1, 0, 0, 0, 2, 0]
]
def keep_only(val, lst):
return [[v if v==val else 0 for v in sublist] for sublist in lst]
ls1 = keep_only(1, ls)
ls2 = keep_only(2, ls)
Output:
print(ls1)
print(ls2)
# [[1, 1, 0, 0, 0], [1, 1, 1, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0]]
# [[0, 0, 0, 2, 2], [0, 0, 0, 2, 2, 2], [0, 0, 0, 0, 0, 2, 0]]
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.