Map a finite probability distribution through a function by summing the masses of source atoms with the same image. The image of the finite support is finite, and fibres partition the original support, so summing the output masses gives the original total one. Non-injective images require addition, not overwriting. A computational snapshot can evaluate each source image once and answer later mass queries by summing matching entries.
Add the following concrete method inside FiniteDist[T]. It forms a pushforward of a finite probability distribution, adding the masses of outcomes that map to the same result:
def map[U](f: T => U): FiniteDist[U] = {
  val source = getDom
  val weighted = source.map(x => (f(x), getProb(x)))
  val target = weighted.map(_._1).distinct
  new FiniteDist[U] {
    def getDom: List[U] = target
    def getProb(y: U): Double =
      weighted.filter(pair => pair._1 == y).map(_._2).sum
  }
}
The input finite probability distribution is assumed to satisfy its interface, with stable masses during construction, and f must be a total deterministic pure function on its finite support. Output values need stable equality consistent with the mathematical equality being modeled; mutating equality-relevant fields afterwards would undermine any finite-support implementation. Computing weighted once also prevents later queries from re-evaluating f inconsistently.
Every source support element appears once in weighted. The distinct operation lists each image value once; every such image has a positive mathematical mass because it has a positive-mass preimage. For each , filtering selects exactly those with , so getProb(y) is the required sum, with an empty sum of zero off the image. The fibers form a set partition of the original support, giving
Thus colliding images have their probabilities added, not overwritten. For a concrete check, mapping a uniform distribution on by parity gives probability to each of ; mapping every source value to one constant gives a singleton domain with total mass one. Floating-point rounding can make a computed total differ slightly from one; the proof is for exact mathematical masses, and the representation uses Double approximations.
Add a generic mapping method inside the supplied trait. Snapshot the finitely many source masses, and sum masses when different entries have the same image:
def map[U](f: T => U): FiniteDist[U] = {
  val entries = getDom.map(x => (f(x), getProb(x)))
  val support = entries.map(_._1).distinct

  new FiniteDist[U] {
    def getDom: List[U] = support
    def getProb(y: U): Double =
      entries.foldLeft(0.0) { case (sum, (image, mass)) =>
        if (image == y) sum + mass else sum
      }
  }
}
Here the mapping is a mathematical function, and keys are assumed to have stable equality consistent with their intended identity. Each source atom is evaluated once when the new distribution is constructed. The distinct image list is exactly the finite support because all listed source atoms have positive mass. Summing collisions, rather than overwriting a previous mass, implements the probability pushforward.
The sets of source atoms with a fixed image partition the original support. Hence nonnegativity is preserved and, in exact probability arithmetic,
Thus the result is a finite probability distribution. For a constant mapping, every atom merges into one outcome of mass one. Floating-point additions in the implementation can introduce rounding error, so numerical tests of normalization should use a suitable tolerance rather than demand exact equality of Double sums.
For a finite probability distribution with support and a total function , assign each image value the combined mass of its preimages. Its support is exactly the finite image , listed without repetitions. Every listed mass is positive because it contains at least one positive source mass.
The fiber of a function construction gives a set partition of , so finite summation gives . This proves nonnegativity and normalization without assuming that is injective. In code, mapping support values and then overwriting repeated image keys is incorrect; repeated keys must have their masses added. For example, mapping a uniform distribution on four consecutive integers to parity produces two outcomes with mass one half each.