puzlink 0.1.0

package/src/lib/distribution.ts

@@ -0,0 +1,162 @@
+import { cumulativeStdNormalProbability as normCdf } from "simple-statistics";
+import { LogCounter } from "./logCounter.js";
+import { LogNum } from "./logNum.js";
+import { memoize } from "./memoize.js";
+import { interval } from "./util.js";
+
+/** A probability distribution of items. */
+export class Distribution<T extends PropertyKey> {
+  private readonly frequencies: ReadonlyMap<T, LogNum>;
+
+  constructor(frequencies: ReadonlyMap<T, LogNum>) {
+    this.frequencies = frequencies;
+  }
+
+  static from(data: string): Distribution<string>;
+  static from<T extends PropertyKey>(data: readonly T[]): Distribution<T>;
+  static from<T extends PropertyKey>(counter: LogCounter<T>): Distribution<T>;
+  static from(data: string | readonly PropertyKey[] | LogCounter<PropertyKey>) {
+    const counter =
+      data instanceof LogCounter ? data : LogCounter.from(data as string);
+    return new Distribution(new Map(counter.frequencies()));
+  }
+
+  /** Get the frequency of the given item. */
+  get(item: T): LogNum {
+    return this.frequencies.get(item) ?? LogNum.from(0);
+  }
+
+  /** Returns an iterable of [item, log probability] pairs. */
+  entries(): IterableIterator<[T, LogNum]> {
+    return this.frequencies.entries();
+  }
+
+  /** k-th moment of the distribution. */
+  @memoize()
+  moment(k: number): LogNum {
+    return LogNum.sum(
+      Array.from(this.frequencies.values(), (freq) => freq.pow(k)),
+    );
+  }
+
+  /** Log probability that k items drawn from the distribution are all equal. */
+  probEqual(k: number): LogNum {
+    if (k <= 0) {
+      return LogNum.from(1);
+    }
+    return this.moment(k);
+  }
+
+  /**
+   * Log probability that k items drawn from the distribution have two distinct
+   * values.
+   */
+  @memoize()
+  probTwoDistinct(k: number): LogNum {
+    const probs = [];
+    for (const i of interval(0, k)) {
+      const j = k - i;
+      probs.push(
+        LogNum.fromBinomial(k, i).mul(this.probEqual(i)).mul(this.probEqual(j)),
+      );
+    }
+    // Case where all are equal is counted 2^k times; others are counted twice:
+    return LogNum.sum(probs)
+      .sub(LogNum.from(2).pow(k).mul(this.probEqual(k)))
+      .div(LogNum.from(2));
+  }
+
+  /**
+   * Log probability that k items drawn from the distribution are all equal,
+   * with exactly one exception.
+   */
+  @memoize()
+  probAlmostEqual(k: number): LogNum {
+    const probs = [];
+    for (const [, freq] of this.frequencies) {
+      probs.push(
+        this.moment(k - 1)
+          .sub(freq.pow(k - 1))
+          .mul(freq)
+          .mul(LogNum.from(k)),
+      );
+    }
+    return LogNum.sum(probs);
+  }
+
+  /** Map the items of the distribution, returning the new distribution. */
+  map<U extends PropertyKey>(fn: (item: T) => U): Distribution<U> {
+    const frequencies = new Map<U, LogNum>();
+    for (const [item, freq] of this.frequencies) {
+      const mapped = fn(item);
+      frequencies.set(
+        mapped,
+        (frequencies.get(mapped) ?? LogNum.from(0)).add(freq),
+      );
+    }
+    return new Distribution(frequencies);
+  }
+
+  /** Chi-squared test statistic against an observed distribution. */
+  chi2(observed: LogCounter<T>): LogNum {
+    for (const [item] of observed.entries()) {
+      if (!this.frequencies.has(item)) {
+        return LogNum.from(Infinity);
+      }
+    }
+    const n = observed.total;
+    const partials = [];
+    for (const [item, freq] of this.frequencies) {
+      const expected = n.mul(freq);
+      const actual = observed.get(item);
+      // Under chi-squared assumptions, the (expected - actual)^2/expected
+      // should be iid N(0, 1)^2.
+      partials.push(expected.absSub(actual).pow(2).div(expected));
+    }
+    return LogNum.sum(partials);
+  }
+
+  /** Log probability that an unordered distribution is drawn from this. */
+  probUnordered(observed: LogCounter<T>): LogNum {
+    const df = this.frequencies.size - 1;
+    const z = (this.chi2(observed).toNum() - df) / Math.sqrt(2 * df);
+    return LogNum.from(1 - normCdf(z));
+  }
+
+  /** Over- and under-represented items, at the given sigma. */
+  outliers(
+    observed: LogCounter<T>,
+    sigma = 2,
+  ): {
+    high: Record<T, LogNum>;
+    low: Record<T, LogNum>;
+  } {
+    const n = observed.total;
+    const low = {} as Record<T, LogNum>;
+    const high = {} as Record<T, LogNum>;
+
+    const keys = [
+      ...this.frequencies.keys(),
+      ...observed.filterKeys((key) => !this.frequencies.has(key)),
+    ];
+    const threshold = LogNum.from(sigma ** 2);
+
+    for (const item of keys) {
+      const freq = this.frequencies.get(item) ?? LogNum.from(0);
+      const expected = n.mul(freq);
+      const actual = observed.get(item);
+
+      // We assume (expected - actual)^2/expected should be distributed as
+      // N(0, 1)^2; thus if it's over sigma^2, it's an outlier.
+      if (expected.absSub(actual).pow(2).div(expected).gt(threshold)) {
+        if (expected.gt(actual)) {
+          low[item] = expected.sub(actual);
+        } else {
+          high[item] = actual.sub(expected);
+        }
+      }
+    }
+
+    return { high, low };
+  }
+}

package/src/lib/lengthDistribution.ts

@@ -0,0 +1,108 @@
+import { Distribution } from "./distribution.js";
+import { LogNum } from "./logNum.js";
+import { memoize } from "./memoize.js";
+import { interval, windows } from "./util.js";
+
+export class LengthDistribution {
+  readonly distribution: Distribution<number>;
+  private readonly distMod2: Distribution<number>;
+  private readonly distMod3: Distribution<number>;
+  /** Length such that 99.9% of answers have length less than this. */
+  private readonly maxLength: number;
+
+  constructor(distribution: Distribution<number>) {
+    this.distribution = distribution;
+    this.distMod2 = this.distribution.map((length) => length % 2);
+    this.distMod3 = this.distribution.map((length) => length % 3);
+
+    this.maxLength = Infinity;
+    let totalProb = LogNum.from(0);
+    for (const [length, freq] of distribution.entries()) {
+      totalProb = totalProb.add(freq);
+      if (totalProb.gt(LogNum.from(0.999))) {
+        this.maxLength = length;
+        break;
+      }
+    }
+  }
+
+  static from(data: Map<number, LogNum>): LengthDistribution {
+    return new LengthDistribution(new Distribution(data));
+  }
+
+  /** Log probability that k words have the same length. */
+  probEqual(k: number): LogNum {
+    return this.distribution.probEqual(k);
+  }
+
+  /** Log probability that k words have the same length, except for one. */
+  probAlmostEqual(k: number): LogNum {
+    return this.distribution.probAlmostEqual(k);
+  }
+
+  /** Log probability that k words have the same length modulo 2. */
+  probEqualMod2(k: number): LogNum {
+    return this.distMod2.probEqual(k);
+  }
+
+  /** Log probability that k words have the same length modulo 3. */
+  probEqualMod3(k: number): LogNum {
+    return this.distMod3.probEqual(k);
+  }
+
+  /** Log probability that k words have consecutive lengths. */
+  @memoize()
+  probConsecutive(k: number): LogNum {
+    if (k <= 1) {
+      return LogNum.from(1);
+    } else if (k > this.maxLength) {
+      return LogNum.from(0);
+    }
+
+    const range = interval(1, this.maxLength).map((i) =>
+      this.distribution.get(i),
+    );
+    const partials = Array.from(windows(range, k), (window) =>
+      LogNum.prod(window),
+    );
+
+    return LogNum.fromFactorial(k).mul(LogNum.sum(partials));
+  }
+
+  /** Log probability that k words have exactly two distinct lengths. */
+  probTwoDistinct(k: number): LogNum {
+    return this.distribution.probTwoDistinct(k);
+  }
+
+  /**
+   * Log probability that k words have distinct lengths, all at least min.
+   */
+  @memoize(2)
+  probDistinct(k: number, min = 0): LogNum {
+    if (k <= 0) {
+      return LogNum.from(1);
+    } else if (min > this.maxLength) {
+      return LogNum.from(0);
+    }
+
+    const probs = [];
+    for (const [length, freq] of this.distribution.entries()) {
+      if (length < min) {
+        continue;
+      }
+      probs.push(freq.mul(this.probDistinct(k - 1, length + 1)));
+    }
+
+    return probs.length === 0
+      ? LogNum.from(0)
+      : LogNum.from(k).mul(LogNum.sum(probs));
+  }
+
+  /** Log probability that k words can be paired by length. */
+  probPaired(k: number): LogNum {
+    if (k % 2 !== 0) {
+      return LogNum.from(0);
+    }
+    return this.probDistinct(Math.floor(k / 2));
+  }
+}

package/src/lib/letterBitset.ts

@@ -0,0 +1,123 @@
+/**
+ * A class that uses a single bigint to store counts for each lowercase letter,
+ * for fast-ish comparison.
+ *
+ * Idea stolen from Collective.jl:
+ * https://github.com/rdeits/Collective.jl/blob/master/src/bitstally.jl
+ */
+export class LetterBitset {
+  private static readonly bits = 5n;
+  private static readonly mask = (1n << LetterBitset.bits) - 1n;
+  private static readonly offsets = Array(26)
+    .fill(0)
+    .map((_, i) => LetterBitset.bits * BigInt(i));
+  private static readonly letterMasks = LetterBitset.offsets.map(
+    (x) => 1n << x,
+  );
+
+  /**
+   * This is a (26 * 5)-bit integer; each 5-bit block is a count for a letter.
+   * Unclear how efficient this'll be for different engines...
+   *
+   * Note that 26 * 5 = 130, so these *mostly* fit in 128-bit integers, unless
+   * there's more than 7 'z's.
+   */
+  readonly data: bigint;
+
+  constructor(data: bigint) {
+    this.data = data;
+  }
+
+  private static toIndex(letter: string) {
+    return letter.charCodeAt(0) - 97;
+  }
+
+  private static fromIndex(index: number) {
+    return String.fromCharCode(97 + index);
+  }
+
+  /** Create a new LetterCounter from a slug. */
+  static from(slug: string) {
+    let data = 0n;
+    for (const char of slug) {
+      data += LetterBitset.letterMasks[this.toIndex(char)]!;
+    }
+    return new LetterBitset(data);
+  }
+
+  /** Count the number of times the given letter appears in this bitset. */
+  index(letter: string): number {
+    return Number(
+      (this.data >> LetterBitset.offsets[LetterBitset.toIndex(letter)]!) &
+        LetterBitset.mask,
+    );
+  }
+
+  equals(other: LetterBitset): boolean {
+    return this.data === other.data;
+  }
+
+  add(char: string): LetterBitset {
+    return new LetterBitset(
+      this.data + (LetterBitset.letterMasks[LetterBitset.toIndex(char)] ?? 0n),
+    );
+  }
+
+  sub(char: string): LetterBitset {
+    return new LetterBitset(
+      this.data - (LetterBitset.letterMasks[LetterBitset.toIndex(char)] ?? 0n),
+    );
+  }
+
+  /** If this + result == other, return result; else null. */
+  transaddOf(other: LetterBitset) {
+    const diff = this.data - other.data;
+    const index = LetterBitset.letterMasks.findIndex((mask) => diff === mask);
+    if (index === -1) {
+      return null;
+    }
+    return LetterBitset.fromIndex(index);
+  }
+
+  /** If this - result == other, return result; else null. */
+  transdeleteOf(other: LetterBitset) {
+    return other.transaddOf(this);
+  }
+}
+
+/** A map from letter bitsets to words with that bitset. */
+export class LetterBitsets {
+  private letterCounters = new Map<bigint, string[]>();
+  private lengths = new Set<number>();
+
+  constructor(wordlist: string[]) {
+    for (const word of wordlist) {
+      const bitset = LetterBitset.from(word).data;
+      if (!this.letterCounters.has(bitset)) {
+        this.letterCounters.set(bitset, []);
+      }
+      this.letterCounters.get(bitset)!.push(word);
+      this.lengths.add(word.length);
+    }
+  }
+
+  /** Get the words whose bitset matches the given slug's bitset. */
+  get(slug: string): string[] {
+    return this.letterCounters.get(LetterBitset.from(slug).data) ?? [];
+  }
+
+  /** Find all substrings of the slug that anagram to a word's bitset. */
+  *matchSubstring(slug: string): Generator<{ start: number; words: string[] }> {
+    for (const length of this.lengths) {
+      let start = 0;
+      let bitset = LetterBitset.from(slug.slice(0, length));
+      for (; start + length <= slug.length; start++) {
+        const words = this.letterCounters.get(bitset.data);
+        if (words && words.length > 0) {
+          yield { start, words };
+        }
+        bitset = bitset.sub(slug[start] ?? "").add(slug[start + length] ?? "");
+      }
+    }
+  }
+}

package/src/lib/letterDistribution.ts

@@ -0,0 +1,236 @@
+import { Distribution } from "./distribution.js";
+import { LogCounter } from "./logCounter.js";
+import { LogNum } from "./logNum.js";
+import { memoize } from "./memoize.js";
+import { caesar, interval } from "./util.js";
+
+export const LETTERS = "abcdefghijklmnopqrstuvwxyz";
+export const VOWELS = "aeiou";
+export const CONSONANTS = "bcdfghjklmnpqrstvwxyz";
+
+/**
+ * Info about the letter distribution of a wordlist.
+ *
+ * All of these methods rely on letters being drawn iid. This is the case for
+ * letters that come from random indexing. This is *not* the case for things
+ * like random substrings.
+ */
+export class LetterDistribution {
+  readonly distribution: Distribution<string>;
+  private readonly lengthToProbs: Map<
+    number,
+    {
+      /** Log probability to get a word of this length. */
+      word: LogNum;
+      /** Log probability to get an anagram of a word of this length. */
+      anagram: LogNum;
+    }
+  >;
+
+  constructor(wordlist: string[]) {
+    const letterCount = new Map<string, number>();
+    let total = 0;
+
+    for (const word of wordlist) {
+      for (const letter of word) {
+        letterCount.set(letter, (letterCount.get(letter) ?? 0) + 1);
+        total++;
+      }
+    }
+
+    this.distribution = new Distribution(
+      new Map(
+        Array.from(letterCount.entries(), ([letter, count]) => [
+          letter,
+          LogNum.fromFraction(count, total),
+        ]),
+      ),
+    );
+
+    this.lengthToProbs = new Map<number, { word: LogNum; anagram: LogNum }>();
+
+    for (const word of wordlist) {
+      const prob = LogNum.prod(
+        Array.from(word, (letter) => this.distribution.get(letter)),
+      );
+
+      const counts = new Map<string, number>();
+      for (const letter of word) {
+        counts.set(letter, (counts.get(letter) ?? 0) + 1);
+      }
+      const perms = LogNum.fromFactorial(word.length).div(
+        LogNum.prod(
+          Array.from(counts.values(), (count) => LogNum.fromFactorial(count)),
+        ),
+      );
+
+      if (!this.lengthToProbs.has(word.length)) {
+        this.lengthToProbs.set(word.length, {
+          word: LogNum.from(0),
+          anagram: LogNum.from(0),
+        });
+      }
+      this.lengthToProbs.get(word.length)!.word = this.lengthToProbs
+        .get(word.length)!
+        .word.add(prob);
+      this.lengthToProbs.get(word.length)!.anagram = this.lengthToProbs
+        .get(word.length)!
+        .anagram.add(prob.mul(perms));
+    }
+  }
+
+  /**
+   * Log probability that a list of letters is iid drawn from this distribution,
+   * via chi-squared.
+   */
+  probUnordered(letters: string[]): LogNum {
+    const counter = LogCounter.from(letters);
+    return this.distribution.probUnordered(counter);
+  }
+
+  /** Over- and under-represented letters, at 2 sigma. */
+  outliers(letters: string[]): {
+    high: string[];
+    low: string[];
+  } {
+    const counter = LogCounter.from(letters);
+    const { high, low } = this.distribution.outliers(counter);
+    return { high: Object.keys(high), low: Object.keys(low) };
+  }
+
+  /**
+   * Log probability that words of given lengths share common letters, in the
+   * same order.
+   */
+  probCommonOrdered(common: number, lengths: number[]): LogNum {
+    if (common === 0) {
+      return LogNum.from(1);
+    }
+
+    const combos = LogNum.prod(
+      interval(0, common - 1).flatMap((k) =>
+        lengths.map((length) => LogNum.fromFraction(length - k, k + 1)),
+      ),
+    );
+    const p = this.distribution.moment(lengths.length).pow(common);
+
+    return LogNum.from(1).sub(LogNum.fromExp(-p.toNum() * combos.toNum()));
+  }
+
+  /** Log probability that k letters are all equal. */
+  probEqual(k: number): LogNum {
+    return this.distribution.probEqual(k);
+  }
+
+  /** Log probability that k letters are all equal, except for one. */
+  probAlmostEqual(k: number): LogNum {
+    return this.distribution.probAlmostEqual(k);
+  }
+
+  /** Log probability that k letters have exactly two values. */
+  probTwoDistinct(k: number): LogNum {
+    return this.distribution.probTwoDistinct(k);
+  }
+
+  /** Log probability that k letters are consecutive. */
+  @memoize()
+  probConsecutive(k: number): LogNum {
+    if (k <= 1) {
+      return LogNum.from(1);
+    } else if (k > LETTERS.length) {
+      return LogNum.from(0);
+    }
+
+    const freqWindow = [];
+    for (const i of interval(1, k)) {
+      freqWindow.push(this.distribution.get(LETTERS[i]!));
+    }
+
+    const partials = [];
+    for (let a = 1; a + k - 1 <= LETTERS.length; a++) {
+      partials.push(LogNum.prod(freqWindow));
+      freqWindow.shift();
+      freqWindow.push(this.distribution.get(LETTERS[a + k]!));
+    }
+
+    return LogNum.fromFactorial(k).mul(LogNum.sum(partials));
+  }
+
+  /**
+   * Log probability that k letters have distinct values, all at least min.
+   */
+  @memoize(2)
+  probDistinct(k: number, min = "a"): LogNum {
+    if (k <= 0) {
+      return LogNum.from(1);
+    } else if (min > "z") {
+      return LogNum.from(0);
+    }
+
+    const probs = [];
+    for (const [letter, freq] of this.distribution.entries()) {
+      if (letter < min) {
+        continue;
+      }
+      probs.push(freq.mul(this.probDistinct(k - 1, caesar(letter, 1))));
+    }
+
+    return probs.length === 0
+      ? LogNum.from(0)
+      : LogNum.from(k).mul(LogNum.sum(probs));
+  }
+
+  /** Log probability that k letters can be grouped in equal pairs. */
+  probPaired(k: number): LogNum {
+    if (k % 2 !== 0) {
+      return LogNum.from(0);
+    }
+    return this.probDistinct(Math.floor(k / 2));
+  }
+
+  /** Log probability that k letters form a word. */
+  probWord(k: number): LogNum {
+    return this.lengthToProbs.get(k)?.word ?? LogNum.from(0);
+  }
+
+  /** Log probability that k letters form an anagram of a word. */
+  probAnagram(k: number): LogNum {
+    return this.lengthToProbs.get(k)?.anagram ?? LogNum.from(0);
+  }
+
+  /** Log probability that k letters are all vowels. */
+  @memoize()
+  probVowels(k: number): LogNum {
+    if (k === 0) {
+      return LogNum.from(1);
+    }
+    if (k === 1) {
+      return LogNum.sum(
+        Array.from(VOWELS, (vowel) => this.distribution.get(vowel)),
+      );
+    }
+    return this.probVowels(1).pow(k);
+  }
+
+  /** Log probability that k letters are all consonants. */
+  @memoize()
+  probConsonants(k: number): LogNum {
+    if (k === 0) {
+      return LogNum.from(1);
+    }
+    if (k === 1) {
+      return LogNum.sum(
+        Array.from(CONSONANTS, (consonant) => this.distribution.get(consonant)),
+      );
+    }
+    return this.probConsonants(1).pow(k);
+  }
+
+  /**
+   * Log probability that n words, each of length k, have the same pattern of
+   * vowels and consonants.
+   */
+  probEqualVowelPattern(n: number, k: number): LogNum {
+    return this.probVowels(k).add(this.probConsonants(k)).pow(n);
+  }
+}

package/src/lib/letterIndices.ts

@@ -0,0 +1,51 @@
+import { enumerate } from "./util.js";
+
+/** A map from letters to their indices in a given slug. */
+export class LetterIndices {
+  private readonly indices: ReadonlyMap<string, number[]>;
+
+  constructor(indices: ReadonlyMap<string, number[]>) {
+    this.indices = indices;
+  }
+
+  static from(slug: string): LetterIndices {
+    const indices = new Map<string, number[]>();
+
+    for (const [i, letter] of enumerate(slug)) {
+      if (!indices.has(letter)) {
+        indices.set(letter, []);
+      }
+      indices.get(letter)!.push(i);
+    }
+
+    return new LetterIndices(indices);
+  }
+
+  *counts(): IterableIterator<[string, number]> {
+    for (const [letter, indices] of this.indices.entries()) {
+      yield [letter, indices.length];
+    }
+  }
+
+  countSet(): Set<number> {
+    return new Set(Array.from(this.counts(), ([, c]) => c));
+  }
+
+  entries(): IterableIterator<[string, number[]]> {
+    return this.indices.entries();
+  }
+
+  filterKeys(fn: (letter: string, indices: number[]) => boolean): string[] {
+    return Array.from(this.indices.entries())
+      .filter(([letter, indices]) => fn(letter, indices))
+      .map(([letter]) => letter);
+  }
+
+  get(letter: string): number[] {
+    return this.indices.get(letter) ?? [];
+  }
+
+  keys(): IterableIterator<string> {
+    return this.indices.keys();
+  }
+}