Pythonic way to “condense” a matrix

Tags: ,



If I have the following matrix, which the input format is a list of lists:

B T E
0 1 0
0 1 1
0 2 1
1 2 0

How can I construct the following python matrix:

       0     1
D = [[{1}, {1,2}],  0
     [{2}, {}]]     1

Where the elements of D, merge the pairs (B,E) with it’s respective T. Example: (0,1) in the above matrix, have T = 1 and T = 2, so in D matrix it should be a set {1,2}. Since there is no (1,1) pair, it should be a empty set {}.

How could a do that in a “pythonic” way?

Answer

You can use collections.defaultdict:

from collections import defaultdict
m = [[0, 1, 0], [0, 1, 1], [0, 2, 1], [1, 2, 0]]
d = defaultdict(dict)
for b, t, e in m:
   d[b][e] = [t] if e not in d[b] else [*d[b][e], t]

l = {i for b in d.values() for i in b}
result = [[set(k.get(j, [])) for j in l] for k in d.values()]
print(result)

Output:

[[{1}, {1, 2}], 
 [{2}, set()]]


Source: stackoverflow