CODEWITHSUNDEEP
pythonpandas
Rick
Rick

Reputation: 45271

pandas multiindex assignment complex broadcasting

I have a dataframe with a multiindex that looks like this:

In [1]: import pandas as pd

In [2]: idx = pd.MultiIndex.from_product([["a", "b"], ["p", "q"], ["x", "y"]])

In [3]: df = pd.DataFrame(index=idx, columns=list("ABCD"))

In [4]: df
Out[4]:
         A    B    C    D
a p x  NaN  NaN  NaN  NaN
    y  NaN  NaN  NaN  NaN
  q x  NaN  NaN  NaN  NaN
    y  NaN  NaN  NaN  NaN
b p x  NaN  NaN  NaN  NaN
    y  NaN  NaN  NaN  NaN
  q x  NaN  NaN  NaN  NaN
    y  NaN  NaN  NaN  NaN

I would like to broadcast an assignment of the value [[1,2,3,4],[5,6,7,8]] to the location "a" at level 0, and ["x", "y"] at level 2, and all locations for level 1. So the resulting dataframe would like this:

         A    B    C    D
a p x    1    2    3    4
    y    5    6    7    8
  q x    1    2    3    4
    y    5    6    7    8
b p x  NaN  NaN  NaN  NaN
    y  NaN  NaN  NaN  NaN
  q x  NaN  NaN  NaN  NaN
    y  NaN  NaN  NaN  NaN

Is this possible in a single assignment operation? My best try has been this, but it doesn't work because the shapes aren't compatible:

In [5]: df.loc["a",:,["x","y"]] = [[1,2,3,4], [5,6,7,8]]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [5], in <module>
----> 1 df.loc["a",:,["x","y"]] = [[1,2,3,4], [5,6,7,8]]
...
ValueError: could not broadcast input array from shape (2,4) into shape (4,4)

The best I've been able to come up with is doing it in two steps:

In [6]: df.loc["a",:,"x"] = [1,2,3,4]

In [7]: df.loc["a",:,"y"] =[5,6,7,8]

Upvotes: 0

Views: 45

Answers (1)

Rick
Rick

Reputation: 45271

Well after posting this I figured it out on my own. Feeling a bit dumb, but here is a way that works fine:

In [16]: df.loc["a"] = [[1,2,3,4], [5,6,7,8]]*2

To make it a bit more general:

In [20]: expand=len(df.index.levels[-2])

In [21]: df.loc["a"] = [[1,2,3,4], [5,6,7,8]]*expand

I'd still like to find a nicer way, but this is ok.

Upvotes: 1

Related Questions