consolidateListsOfObjects Typescript

consolidate.ts

export const consolidateListsOfObjects = <Obj>(args: {
  readonly objectLists: ReadonlyArray<ReadonlyArray<Obj>>;
  readonly objectKeyFn: (obj: Obj) => string;
}): ReadonlyArray<ReadonlyArray<Obj>> => {
  const parentMap = new Map<string, string>();
  const find = (objKey: string): string => {
    const parent = parentMap.get(objKey);
    if (parent === undefined) {
      return objKey;
    }
    const newParent = find(parent);
    parentMap.set(objKey, newParent);
    return newParent;
  };
  const union = (objKey1: string, objKey2: string): void => {
    const parent1 = find(objKey1);
    const parent2 = find(objKey2);
    if (parent1 !== parent2) {
      parentMap.set(parent2, parent1);
    }
  };

  const { objectLists, objectKeyFn } = args;
  objectLists.forEach((objList) => {
    objList.forEach((obj, idx) => {
      if (idx > 0) {
        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
        union(objectKeyFn(objList[idx - 1]!), objectKeyFn(obj));
      }
    });
  });

  const parentToConsolidatedListMap = new Map<string, Array<Obj>>();
  const seenKeys = new Set<string>();
  objectLists.forEach((objList) => {
    objList.forEach((obj) => {
      const objKey = objectKeyFn(obj);
      if (!seenKeys.has(objKey)) {
        seenKeys.add(objKey);
        const parent = find(objKey);
        if (!parentToConsolidatedListMap.has(parent)) {
          parentToConsolidatedListMap.set(parent, []);
        }
        parentToConsolidatedListMap.get(parent)?.push(obj);
      }
    });
  });

  return [...parentToConsolidatedListMap.values()];
};