Finding a row in a 2d array in python if the value of the column is known

Viewed 59

I have a matrix in an excel sheet I am reading into my script using Pandas. I convert it to an np matrix like so and come out with this as a result.

df = pd.read_excel(r'C:\Users\PycharmProjects\OLS_Script\ols1.xlsx')
matrix = np.matrix(matrix)
print(matrix)

enter image description here

Now my question is regarding if some mechanics in python exist that I am not familiar with as I come from C++. Essentially what I am trying to do with this matrix is for each column, acquire the entire row that the column has values in, as well as the row corresponding to that.

For instance, looking at column 0, I would need the entire row 0, as well as row 1 since it is the only other row with a 1 in column 0. So I would need the indices [1,2,3] in row 0, and [4,5,6,7] in row 1, excluding anything in column 0.

Column 1 has a 1 in row 0 as well as row 2, so I would need to get those two rows and the corresponding column indices in each of those rows with a population in it. [0,2,3] in row 0, [4,8,9,10,11] in row 2, excluding anything in column 1, and so on for the following columns.

My original idea was the have a for loop go through each column, then another for loop inside to go through each row at the same column index to find where the next corresponding row is. I am not too familiar with python and the various "shortcuts" it has compared to something like C++ since I have seen a lot of functionality in python beforehand that might make something like what I need to do with as few lines of code as possible, so if there is a shorter way to accomplish what I need, please let me know.

1 Answers

With an excerpt of your matrix as an example:

import pandas as pd

matrix = [
    [1, 1, 1, 1, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA,],
    [1, pd.NA, pd.NA, pd.NA, 1, 1, 1, 1, pd.NA, pd.NA, pd.NA, pd.NA,],
    [pd.NA, 1, pd.NA, pd.NA, 1, pd.NA, pd.NA, pd.NA, 1, 1, 1, 1],
]

df = pd.DataFrame(matrix)
     0     1     2     3     4     5     6     7     8     9     10    11
0     1     1     1     1  <NA>  <NA>  <NA>  <NA>  <NA>  <NA>  <NA>  <NA>
1     1  <NA>  <NA>  <NA>     1     1     1     1  <NA>  <NA>  <NA>  <NA>
2  <NA>     1  <NA>  <NA>     1  <NA>  <NA>  <NA>     1     1     1     1

Here is one way to do it:

results = {
    i: df.dropna(subset=i).dropna(how="all", axis=1).drop(columns=i).columns.to_list()
    for i in range(df.shape[1])
}
print(results)
# Output
{
    0: [1, 2, 3, 4, 5, 6, 7],
    1: [0, 2, 3, 4, 8, 9, 10, 11],
    2: [0, 1, 3],
    3: [0, 1, 2],
    4: [0, 1, 5, 6, 7, 8, 9, 10, 11],
    5: [0, 4, 6, 7],
    6: [0, 4, 5, 7],
    7: [0, 4, 5, 6],
    8: [1, 4, 9, 10, 11],
    9: [1, 4, 8, 10, 11],
    10: [1, 4, 8, 9, 11],
    11: [1, 4, 8, 9, 10],
}
Related