in internal/stats/udist.go [165:297]
func makeUmemo(twoU, n1 int, t []int) []map[ukey]float64 {
// Another candidate for a fast implementation is van de Wiel,
// "The split-up algorithm: a fast symbolic method for
// computing p-values of distribution-free statistics". This
// is what's used by R's coin package. It's a comparatively
// recent publication, so it's presumably faster (or perhaps
// just more general) than previous techniques, but I can't
// get my hands on the paper.
//
// TODO: ~40% of this function's time is spent in mapassign on
// the assignment lines in the two loops and another ~20% in
// map access and iteration. Improving map behavior or
// replacing the maps altogether with some other constant-time
// structure could double performance.
//
// TODO: The worst case for this function is when there are
// few ties. Yet the best case overall is when there are *no*
// ties. Can we get the best of both worlds? Use the fast
// algorithm for the most part when there are few ties and mix
// in the general algorithm just where we need it? That's
// certainly possible for sub-problems where t[:k] has no
// ties, but that doesn't help if t[0] has a tie but nothing
// else does. Is it possible to rearrange the ranks without
// messing up our computation of the U statistic for
// sub-problems?
K := len(t)
// Compute a coefficients. The a slice is indexed by k (a[0]
// is unused).
a := make([]int, K+1)
a[1] = t[0]
for k := 2; k <= K; k++ {
a[k] = a[k-1] + t[k-2] + t[k-1]
}
// Create the memo table for the counts function, A. The A
// slice is indexed by k (A[0] is unused).
//
// In "The Mann Whitney Distribution Using Linked Lists", they
// use linked lists (*gasp*) for this, but within each K it's
// really just a memoization table, so it's faster to use a
// map. The outer structure is a slice indexed by k because we
// need to find all memo entries with certain values of k.
//
// TODO: The n1 and twoU values in the ukeys follow strict
// patterns. For each K value, the n1 values are every integer
// between two bounds. For each (K, n1) value, the twoU values
// are every integer multiple of a certain base between two
// bounds. It might be worth turning these into directly
// indexible slices.
A := make([]map[ukey]float64, K+1)
A[K] = map[ukey]float64{ukey{n1: n1, twoU: twoU}: 0}
// Compute memo table (k, n1, twoU) triples from high K values
// to low K values. This drives the recurrence relation
// downward to figure out all of the needed argument triples.
//
// TODO: Is it possible to generate this table bottom-up? If
// so, this could be a pure dynamic programming algorithm and
// we could discard the K dimension. We could at least store
// the inputs in a more compact representation that replaces
// the twoU dimension with an interval and a step size (as
// suggested by Cheung, Klotz, not that they make it at all
// clear *why* they're suggesting this).
tsum := sumint(t) // always ∑ t[0:k]
for k := K - 1; k >= 2; k-- {
tsum -= t[k]
A[k] = make(map[ukey]float64)
// Construct A[k] from A[k+1].
for A_kplus1 := range A[k+1] {
rkLow := maxint(0, A_kplus1.n1-tsum)
rkHigh := minint(A_kplus1.n1, t[k])
for rk := rkLow; rk <= rkHigh; rk++ {
twoU_k := A_kplus1.twoU - rk*(a[k+1]-2*A_kplus1.n1+rk)
n1_k := A_kplus1.n1 - rk
if twoUmin(n1_k, t[:k], a) <= twoU_k && twoU_k <= twoUmax(n1_k, t[:k], a) {
key := ukey{n1: n1_k, twoU: twoU_k}
A[k][key] = 0
}
}
}
}
// Fill counts in memo table from low K values to high K
// values. This unwinds the recurrence relation.
// Start with K==2 base case.
//
// TODO: Later computations depend on these, but these don't
// depend on anything (including each other), so if K==2, we
// can skip the memo table altogether.
if K < 2 {
panic("K < 2")
}
N_2 := t[0] + t[1]
for A_2i := range A[2] {
Asum := 0.0
r2Low := maxint(0, A_2i.n1-t[0])
r2High := (A_2i.twoU - A_2i.n1*(t[0]-A_2i.n1)) / N_2
for r2 := r2Low; r2 <= r2High; r2++ {
Asum += mathChoose(t[0], A_2i.n1-r2) *
mathChoose(t[1], r2)
}
A[2][A_2i] = Asum
}
// Derive counts for the rest of the memo table.
tsum = t[0] // always ∑ t[0:k-1]
for k := 3; k <= K; k++ {
tsum += t[k-2]
// Compute A[k] counts from A[k-1] counts.
for A_ki := range A[k] {
Asum := 0.0
rkLow := maxint(0, A_ki.n1-tsum)
rkHigh := minint(A_ki.n1, t[k-1])
for rk := rkLow; rk <= rkHigh; rk++ {
twoU_kminus1 := A_ki.twoU - rk*(a[k]-2*A_ki.n1+rk)
n1_kminus1 := A_ki.n1 - rk
x, ok := A[k-1][ukey{n1: n1_kminus1, twoU: twoU_kminus1}]
if !ok && twoUmax(n1_kminus1, t[:k-1], a) < twoU_kminus1 {
x = mathChoose(tsum, n1_kminus1)
}
Asum += x * mathChoose(t[k-1], rk)
}
A[k][A_ki] = Asum
}
}
return A
}