Expectation of piecewise constant function

Question

Say I have a set $\mathcal{X}$ and a partition $A := (A_{1}, \dots, A_{M})$, i.e. subsets such that $\dot{\cup}_{i=1}^M A_{i} = \mathcal{X}$ and the $A_{i}$s are pairwise disjoint. Say further I have a function $f: \mathcal{X} \to \mathcal{Y}$ that is piecewise constant on the cells of $A$, i.e.

$$ f (x) = \begin{cases} f_{1} & x \in A_{1} \\ \dots \\ f_{M} & x \in A_{M} \end{cases} $$

Consider now a random variable $X$ that is evenly distributed on $\mathcal{X}$, i.e. its probability density function is $p(x) = \varphi ~ ,\forall x$.

I am interested in expressing the expectation of $f(x)$ over $\mathcal{X}$ in terms of the function values $f_{1}, \dots, f_{m}$ at the cells of the partition.

Here's what I have so far but I am not sure it's correct. Mostly, I am unsure whether, in the last line, the individual summands should be having a weight factor that corresponds to the sizes of the cells. $$ \begin{align} \mathbb{E}_{X}\left[ f(X) \right] &= \int f(x) \varphi \, dx \\ &= \int_{A_{1}} f_{1} \varphi \, dx ~ ~ + \dots + \int_{A_{M}} f_{M}\varphi \, dx \end{align} $$ Now let $X_{|i}$ be a random variable that has pdf of $X$ in cell $A_{i}$ and $0$ everywhere else. Then $$ \begin{align} \dots &= \mathbb{E}_{X_{|1}}\left[ f_{1} \right] ~ ~ + \dots + \mathbb{E}_{X_{|M}}\left[ f_{M} \right] \end{align} $$ and, since the expectation over a constant is just that constant, we'd have $$ \begin{align} \dots &= f_{1} + \dots + f_{M} \end{align} $$ This seems weird since the expectation $\mathbb{E}_{X}\left[ f(X) \right]$ would then be the sum of the individual values whereas intuitively I'd expect it to be more like a weighted mean.

Answer 1

As @J. Delaney noticed you didn't normalize it properly.

$$ \Pr(X \in A_i) = \int_{A_i} p(x) dx $$

so

$$ \int_{A_i} f(x) p(x) dx $$

is not the expected value as it integrates over just a part of the support for $x$. So you are not summing the expected values. The integration is correct, you just misunderstood the integrals over the $A_i$ subsets as expected values.

If you had conditional expectations

$$ E[f(X) | X \in A_i] = \int f(x) \frac{P(X = x, X \in A_i)}{P(X \in A_i)} dx $$

then indeed you would have to use the weighted average of those with the weights equal to $P(X \in A_i)$ to cancel the denominator above (see also the law of total expectation).