Home > Net >  Python: create pivot table as numpy 2d-array from two lists
Python: create pivot table as numpy 2d-array from two lists

Time:08-05

I have the following lists.

list_1 = ["A", "B", "C"]
list_2 = ["D", "A", "B"]

And I want to create a numpy 2d-array that would work as a pivot table. The columns should be values from list_1, rows should be values from list_2, and values should be booleans saying if the values are/are not equal. The result should look like this.

array([[0., 0., 0.],
       [1., 0., 0.],
       [0., 1., 0.]])

I have managed to write the code on my own, but I don't think my solution is the elegant one. I hope there is some built in numpy function for this that would make things simpler. Can you help me with that, please?

My code is following.

import numpy as np

list_1 = ["A", "B", "C"]
list_2 = ["D", "A", "B"]

my_array = np.zeros((len(list_2), len(list_1)))

for i, val in enumerate(list_2):
    my_array[i] = val == np.array(list_1)

print(my_array)

CodePudding user response:

Use numpy broadcasting:

list_1 = ["A", "B", "C"]
list_2 = ["D", "A", "B"]

my_array = (np.array(list_1) == np.array(list_2)[:,None]).astype(int)

output:

array([[0, 0, 0],
       [1, 0, 0],
       [0, 1, 0]])

NB. if you really want booleans (True/False), leave out the .astype(int):

array([[False, False, False],
       [ True, False, False],
       [False,  True, False]])
  • Related