Groups at a party

groups_at_a_party.ts

function estimateExpectedNumberOfGroups(n: number, iterations: number = 2000): number {
  const groupCounts = range(iterations).map(() =>
    countGroups(createPairsOfPeople(n))
  );
  return mean(groupCounts);
}

type Pair<T> = readonly [T, T];

function countGroups(pairsOfPeople: Iterable<Pair<number>>): number {
  const relationships = DirectedGraph.fromEdges(pairsOfPeople);
  const seen = new Set()
  let groups = 0;

  for (const [person] of pairsOfPeople) {
    if (seen.has(person)) {
      continue;
    }

    const peopleInSameGroup = relationships.walkDfs(person);
    for (const otherPerson of peopleInSameGroup) seen.add(otherPerson);

    groups += 1;
  }

  return groups;
}

function createPairsOfPeople(n: number): Map<number, number> {
  const remainingPeople = new Set(range(n));
  const pairsOfPeople = new Map();
  for (let personA = 0; personA < n; personA++) {
    const personB = selectRandomItem([...remainingPeople]);
    remainingPeople.delete(personB);

    pairsOfPeople.set(personA, personB);
  }
  return pairsOfPeople;
}

class DirectedGraph<T> {
  private readonly adjacentNodes: Map<T, Set<T>> = new Map();

  static fromEdges<T>(edges:  Iterable<Pair<T>>): DirectedGraph<T> {
    const digraph = new DirectedGraph<T>();
    for (const [head, tail] of edges) digraph.addEdge(head, tail);
    return digraph;
  }

  addEdge(head: T, tail: T) {
    let tails = this.adjacentNodes.get(head);
    if (tails == null) {
      tails = new Set();
      this.adjacentNodes.set(head, tails);
    }

    tails.add(tail);
  }

  walkDfs(source: T): Set<T> {
    const visitedNodes = this.dfsIter(source, new Set<T>());
    return new Set(visitedNodes);
  }

  private *dfsIter(entryPointNode: T, previouslyVisitedNodes: Set<T>): IterableIterator<T> {
    if (previouslyVisitedNodes.has(entryPointNode)) {
      return;
    }

    previouslyVisitedNodes.add(entryPointNode);
    yield entryPointNode;

    if (this.adjacentNodes.has(entryPointNode)) {
      const tails = this.adjacentNodes.get(entryPointNode)!;

      for (const tail of tails) {
        yield* this.dfsIter(tail, previouslyVisitedNodes);
      }
    }
  }
}

function range(length: number): number[] {
  const nums = [];
  for (let i = 0; i < length; i++) nums.push(i);
  return nums;
}

function selectRandomItem<T>(values: readonly T[]): T {
  return values[randomInt(values.length)];
}

function randomInt(n: number): number {
  return Math.floor(Math.random() * n);
}

function mean(values: readonly number[]): number {
  let sum = 0
  for (const value of values) sum += value
  return sum / values.length
}