Create xarray Dataset with observations and averages that has combined index

Suppose I have the following dataarray containing observations for different locations over time:

import numpy as np
import pandas as pd
import xarray as xr

np.random.seed(42)

data = xr.DataArray(
    np.random.randint(1,100, (36, 3)), 
    dims=("time", "location"), 
    coords={
        "time": pd.date_range("2022-01-01", periods=36, freq="10D"), 
        "location": ["A", "B", "C"]
    },
    name="observations"
)

and now I calculate the monthly average and combine it with the observations to a dataset:

monthly_avg = data.groupby("time.month").mean()
data = data.to_dataset()
data["average"] = monthly_avg

giving me

How can is set the indices correctly (if possible) so when I run:

data.sel(time="2022-01-01")

I get a subset of the dataset for one time, all locations and one monthly average (which corresponds to the selected time)?

At the moment when I run this I get

returning all monthly averages for the timestep.

Conversely, when I run

data.sel(month=1)

I'd like a subset with only the timesteps that are in January.

To get the selection return what you want, I would first compute the monthly averages and repeat them to match the original time-dimension.
Then I would create a multi-index, such that you can select either the specific date or the month.

#setup test data 
import numpy as np
import pandas as pd
import xarray as xr

np.random.seed(42)

data = xr.DataArray(
    np.random.randint(1,100, (36, 3)), 
    dims=("time", "location"), 
    coords={
        "time": pd.date_range("2022-01-01", periods=36, freq="10D"), 
        "location": ["A", "B", "C"]
    },
    name="observations"
)

#compute monthly array and repeat with list comprehension
data=data.to_dataset()
monthly_avg = data.groupby("time.month").mean()['observations'].values
data['average']=(('time','location'),np.array([monthly_avg[i-1,:] for i in data.time.dt.month]))

#add month  and create multiindex
data['month']=data.time.dt.month
data=data.set_index(day_month=['time','month'])

You can then run the selection to get what you want.

print(data.sel(time="2022-01-01"))
<xarray.Dataset>
Dimensions:       (location: 3, month: 1)
Coordinates:
  * location      (location) <U1 'A' 'B' 'C'
  * month         (month) int64 1
    time          <U10 '2022-01-01'
Data variables:
    observations  (month, location) int64 52 93 15
    average       (month, location) float64 70.5 82.25 33.75

print(data.sel(month=1))
<xarray.Dataset>
Dimensions:       (location: 3, time: 4)
Coordinates:
  * location      (location) <U1 'A' 'B' 'C'
  * time          (time) datetime64[ns] 2022-01-01 2022-01-11 ... 2022-01-31
    month         int64 1
Data variables:
    observations  (time, location) int64 52 93 15 72 61 21 83 87 75 75 88 24
    average       (time, location) float64 70.5 82.25 33.75 ... 70.5 82.25 33.7

This gives repeated values for the second command.

Maybe there is a better way to set up the multi-index.
You can have a look at the pandas multiindex documentation: https://pandas.pydata.org/docs/user_guide/advanced.html:
or into stack/unstack in xarray: https://xarray.pydata.org/en/v0.7.2/reshaping.html#stack-and-unstack
in case you haven't done so before.