CODEWITHSUNDEEP
pythonmatrix
Egydio Pacheco
Egydio Pacheco

Reputation: 103

Pythonic way to "condense" a matrix

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?

Upvotes: 3

Views: 65

Answers (2)

Ajax1234
Ajax1234

Reputation: 71461

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()]]

Upvotes: 4

Sayandip Dutta
Sayandip Dutta

Reputation: 15872

It's hard to guess what your input or output format is. If you are using pandas, you can do:

>>> data
[[0, 1, 0],
 [0, 1, 1],
 [0, 2, 1],
 [1, 2, 0]]

>>> df = pd.DataFrame(data, columns=['B', 'T', 'E'])
>>> df
   B  T  E
0  0  1  0
1  0  1  1
2  0  2  1
3  1  2  0

>>> df.groupby(['B', 'E']).agg(set).unstack('E', fill_value=set())
     T        
E    0       1
B             
0  {1}  {1, 2}
1  {2}      {}

# OR,
>>> df.groupby(['B', 'E']).agg(set).unstack('E', fill_value=set()).to_numpy()

array([[{1}, {1, 2}],
       [{2}, set()]], dtype=object)

Upvotes: 2

Related Questions