ustat_exact.js

// Copyright 2019 Kris Selden BSD 2-CLAUSE LICENSE

// math used is public domain just dont copy my code without giving credit.

// @ts-check

/**
 * @typedef {{get(U: number, n1: number, n2: number): number; set(U: number, n1: number, n2: number, count: number): void;}} ICache
 */

/**
 * Count frequency of U (rank sum - min rank sum)
 * @param U U statistic which is the rank sum - min rank sum (W - (n1 * (n2 + 1)))
 * @param n1 {number} sample count from set 1
 * @param n2 {number} sample count from set 2
 * @param cache {ICache}
 * @returns {number} the frequency of U
 */
function countU(U, n1, n2, cache = newCache()) {
  let max = n1 * n2;

  if (U < 0 || n1 < 0 || n2 < 0 || U > max) {
    return 0;
  }

  // count is symetrical around the midpoint
  if (U > max >> 1) {
    U = max - U;
  }

  if (n2 < n1) {
    // count of U in 3, 5 is same as 5, 3
    // the smaller side has less recursion
    return countU(U, n2, n1, cache);
  }

  if (U === 0) {
    return 1;
  }

  if (n1 === 0 || n2 === 0) {
    return 0;
  }

  if (U < n2) {
    // fu(u - n2, n1 - 1, n2) + fu(u, n1, n2 - 1)
    // until u - n2 === 0 only the right branch is taken fu(u, n1, n2 - 1)
    // so jump ahead to that result
    return countU(U, n1, U, cache);
  }

  // Recursive formula
  // FU(u|n1,n2) = FU(u−n2|n1 −1,n2) + FU(u|n1,n2 −1)

  let count = cache.get(U, n1, n2);
  if (count === undefined) {
    count = countU(U - n2, n1 - 1, n2, cache) + countU(U, n1, n2 - 1, cache);
    cache.set(U, n1, n2, count);
  }

  return count;
}

function newCache() {
  // assigning numbers to an object will be backed by a NumberDictionary
  const cache = {};
  return {
    cache,

    /**
     * @param {number} U
     * @param {number} n1
     * @param {number} n2
     * @returns {number | undefined}
     */
    get(U, n1, n2) {
      const key = (n1 << 24) | (n2 << 16) | U;
      const count = cache[key];
      if (count === undefined) {
        this.misses++;
      } else {
        this.hits++;
      }
      return count;
    },

    /**
     * @param {number} U
     * @param {number} n1
     * @param {number} n2
     * @param {number} count
     */
    set(U, n1, n2, count) {
      const key = (n1 << 24) | (n2 << 16) | U;
      cache[key] = count;
    },
    misses: 0,
    hits: 0
  };
}

/**
 * @param {number} from
 * @param {number} to
 * @param {number=} by
 */
function* sequence(from, to, by = 1) {
  for (let i = from; i <= to; i += by) yield i;
}

/**
 * @param {Iterable<number>} x
 */
function* cumulativeSum(x) {
  let sum = 0;
  for (const y of x) {
    sum += y;
    yield sum;
  }
}

/**
 * @param {Iterable<number>} x
 * @param {(number) => number} fn
 */
function* map(x, fn) {
  for (const y of x) {
    yield fn(y);
  }
}

/**
 * @param {number} n1
 * @param {number} n2
 * @param {ICache} cache
 */
function* freq(n1, n2, cache) {
  yield* cumulativeSum(
    map(sequence(0, n1 * n2), U => countU(U, n1, n2, cache))
  );
}

let cache = newCache();

console.log(new Int32Array(freq(3, 5, cache)));
console.log(new Int32Array(freq(4, 5, cache)));
console.log(new Int32Array(freq(5, 5, cache)));
console.log(new Int32Array(freq(10, 10, cache)));

console.log(new Int32Array(freq(20, 20, cache)));

cache = newCache();
console.time("U 50 50");
console.log(new Int32Array(freq(50, 50, cache)));
console.timeEnd("U 50 50");

cache = newCache();
console.time("U 50 50");
const m = 50;
const n = 50;
const max = m * n;
const results = new Int32Array(max + 1);
let sum = 0;
for (let U = 0; U <= max; U++) {
  sum += countU(U, m, n, cache);
  results[U] = sum;
}
console.timeEnd("U 50 50");
console.log(results);