bestGroupBy.ts

bestGroupBy.ts

/**
 * Like "groupBy", but only track the "best" group.
 * @param arr Array of values
 * @param grouper Mapper from elements of array to some group. Should be pure (it'll be called on first element twice).
 * @param compare Compare two groups: if `compare(x, y) > 0`, then `x > y`, and similarly `=` and `<`.
 * @return array containing the best group as first element, then the array of entries in that group.
 * 
 * I only care about `compare`'s output's sign/zero, so it doesn't need to be a true distance.
 */
export function bestGroupBy<T, U>(arr: T[], grouper: (x: T) => U, compare: (x: U, y: U) => number): [U, T[]] {
  let bestGroup: U = grouper(arr[0]);
  let bestTs: T[] = [];
  for (const x of arr) {
    const groupx = grouper(x);
    const compared = compare(groupx, bestGroup);
    if (compared === 0) {
      bestTs.push(x);
    } else if (compared > 0) {
      bestTs = [x];
      bestGroup = groupx;
    }
  }
  return [bestGroup, bestTs];
}

groupBy.ts

/* A single function that does both `groupBy` and `bestGroupBy`, though it's a bit unnatural to have this `compare` argument here. */

export function flexibleGroupBy<T, U>(arr: T[], grouper: (x: T) => U, compare?: (x: U, y: U) => number): Map<U, T[]> {
  let init: Map<U, T[]> = new Map([[grouper(arr[0]), []]]);
  return arr.reduce((memo, x) => {
    const group = grouper(x);
    if (memo.has(group)) {
      (memo.get(group) || []).push(x);
      return memo;
    }
    if (compare && compare(group, Array.from(memo.keys())[0]) > 0) { return new Map([[group, [x]]]) }
    memo.set(group, [x]);
    return memo;
  }, init);
}
export const groupBy = flexibleGroupBy;
export const bestGroupBy = flexibleGroupBy;