CODEWITHSUNDEEP
pythonpandasnumpy
mgalardini
mgalardini

Reputation: 1467

Stack a square DataFrame to only keep the upper/lower triangle

I have a symmetric square DataFrame in pandas:

a = np.random.rand(3, 3)
a = (a + a.T)/2
np.fill_diagonal(a, 1.)
a = pd.DataFrame(a)

That looks like this:

          0         1         2
0  1.000000  0.747064  0.357616
1  0.747064  1.000000  0.631622
2  0.357616  0.631622  1.000000

If I apply the stack method, I would get lots of redundant information (including the diagonal, in which I'm not interested):

0  0    1.000000
   1    0.747064
   2    0.357616
1  0    0.747064
   1    1.000000
   2    0.631622
2  0    0.357616
   1    0.631622
   2    1.000000

Is there a way to only get the lower (or upper) triangle this using "pure" pandas?

1  0    0.747064
2  0    0.357616
   1    0.631622

Upvotes: 7

Views: 3849

Answers (3)

Soudipta Dutta
Soudipta Dutta

Reputation: 2122

import numpy as np
import pandas as pd

data = {
    0: [100, 200, 300],
    1: [400, 500, 600],
    2: [700, 800, 1000]
}

a = pd.DataFrame(data)

# Create a mask for the upper triangle
mask = np.triu(np.ones_like(a, dtype=bool), k=1)
'''
[[False  True  True]
 [False False  True]
 [False False False]]
'''
a = a.where(mask).stack()
print(a)
'''
0  1    400.0
   2    700.0
1  2    800.0
dtype: float64
'''

Upvotes: 0

Zero
Zero

Reputation: 76917

You could use mask

In [278]: a.mask(np.triu(np.ones(a.shape)).astype(bool)).stack()
Out[278]:
1  0    0.747064
2  0    0.357616
   1    0.631622
dtype: float64

Or use where

In [285]: a.where(np.tril(np.ones(a.shape), -1).astype(bool)).stack()
Out[285]:
1  0    0.747064
2  0    0.357616
   1    0.631622
dtype: float64

Upvotes: 8

mgalardini
mgalardini

Reputation: 1467

The easiest way I could think of is to force the upper (or lower) triangle to NaN, as by default the stack method will not include NaNs:

a.values[np.triu_indices_from(a, 0)] = np.nan
a.stack()

which gives:

1  0    0.747064
2  0    0.357616
   1    0.631622

Upvotes: 6

Related Questions