Search code examples
pythonpandasnumpydataframediagonal

Align dataframe diagonals into columns?


consider the pd.DataFrame df

df = pd.DataFrame([
        [1, 2, 3, 4, 5],
        [5, 1, 2, 3, 4],
        [4, 5, 1, 2, 3],
        [3, 4, 5, 1, 2],
        [2, 3, 4, 5, 1]
    ], list('abcde'), list('ABCDE'))

How do I align the diagonal values into columns?

I'd like this as the result

enter image description here


i've done this

pd.DataFrame([np.roll(row, -k) for k, (_, row) in enumerate(df.iterrows())],
             df.index, df.columns)

I'm hoping for something more straight forward.


Solution

  • You can use numpy solution - for shift is used reversed Series same length as DataFrame (if DataFrame has non numeric and non monotonic index it works nice also):

    A = df.values
    r = pd.Series(range(len(df)))[::-1] + 1
    
    rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]
    
    r[r < 0] += A.shape[1]
    column_indices = column_indices - r[:,np.newaxis]
    
    result = A[rows, column_indices]
    print (pd.DataFrame(result, df.index, df.columns))
       A  B  C  D  E
    a  1  2  3  4  5
    b  1  2  3  4  5
    c  1  2  3  4  5
    d  1  2  3  4  5
    e  1  2  3  4  5