I have a problem with splitting a nested list. There is an example of my input:

l = [[[35, 58, 'A'], [0, 18, 'B'], [76, 101, 'B'], [103, 130, 'A'], [134, 158, 'A']], [[2, 51, 'A'], [55, 115, 'B'], [125, 150, 'B']]]

I want to split it into two separate nested lists by the third element:

l_b = [[[0,18], [76,101]], [[55, 115], [125,150]]]
l_a = [[[35,58], [103, 130], [134,158]], [[2,51]]]

I tried to search by using a loop for:

for i in l:
    for ii in i:
        if ii[2] == 'A':
            l_a.append(ii)

But I simply get a flat list, so I am losing the indexes. Thanks in advance ;)

CodePudding user response：

Unnest the list, iterate it and insert a slice of each element into a dict under a key.

l = [[[35, 58, 'A'], [0, 18, 'B'], [76, 101, 'B'], [103, 130, 'A'], 
     [134, 158, 'A']], [[2, 51, 'A'], [55, 115, 'B'], [125, 150, 'B']]]

flat = [k for x in l for k in x]
print(flat)

from collections import defaultdict

# this will automatically create the key with a list if needed
# no need to prepopulate it if you got "C"s and "D"s later on as well
grp = defaultdict(list)

for elem in flat:
    # use 3rd element as key, use slice of first two as data
    # will crash on less then 3 elements but your data is ok
    grp[elem[2]].append(elem[:2])

print(*grp.items(), sep="\n")

Output:

#flat
[[35, 58, 'A'], [0, 18, 'B'], [76, 101, 'B'], [103, 130, 'A'], 
 [134, 158, 'A'], [2, 51, 'A'], [55, 115, 'B'], [125, 150, 'B']]

# grouped
('A', [[35, 58], [103, 130], [134, 158], [2, 51]])
('B', [[0, 18], [76, 101], [55, 115], [125, 150]])

If needed you can re-nest the list - just see not much sense in it:

l_a = [grp["A"]]

You could simple use l_a = [l_a] after your code to renest it.

CodePudding user response：

If you need to keep the nested structure for some reason:

l = [[[35, 58, 'A'], [0, 18, 'B'], [76, 101, 'B'], [103, 130, 'A'], [134, 158, 'A']], [[2, 51, 'A'], [55, 115, 'B'], [125, 150, 'B']]]
l_a = []

for i in l:
    tmp_list = []
    for ii in i:
        if ii[2] == 'A':
            tmp_list.append(ii[:2])
    l_a.append(tmp_list)

CodePudding user response：

First create empty lists l_a, and l_b, after the if statement, you need to slice ii[:2] to include the numbers only. Lastly, include else statement in order to append the values where B exists.

CodePudding user response：

l = [[[35, 58, 'A'], [0, 18, 'B'], [76, 101, 'B'], [103, 130, 'A'], [134, 158, 'A']], [[2, 51, 'A'], [55, 115, 'B'], [125, 150, 'B']]]
l_a_outer =[]
for i in l:
    l_a = []
    for ii in i:
        if ii[2] == 'A':
            l_a.append(ii)
    l_a_outer.append(l_a)
print(l_a_outer)

sample output

[[[35, 58, 'A'], [103, 130, 'A'], [134, 158, 'A']], [[2, 51, 'A']]

CodePudding user response：

As you target "A" and "B" specifically, you can prepare a dict just for those two lists, and then populate those lists via the dict keys:

l = [[[35, 58, 'A'], [0, 18, 'B'], [76, 101, 'B'], [103, 130, 'A'], [134, 158, 'A']], [[2, 51, 'A'], [55, 115, 'B'], [125, 150, 'B']]]

l_a = []
l_b = []
res = { "A": l_a, "B": l_b }
for chunk in l:
    l_a.append([])
    l_b.append([])
    for *values, letter in chunk:
        res[letter][-1].append(values)

print(l_a)  # [[[35, 58], [103, 130], [134, 158]], [[2, 51]]]
print(l_b)  # [[[0, 18], [76, 101]], [[55, 115], [125, 150]]]