✨ feat(@algorithm.ts/findset): support EnhancedFindset (use through )

guanghechen · guanghechen · commit da80941d186c · 2021-11-28T14:35:28.000+08:00
diff --git a/packages/findset/README-zh.md b/packages/findset/README-zh.md
@@ -127,6 +127,22 @@
   findset.size(4)   // => 1
   ```
 
+* 创建一个增强的并查集
+
+  ```typescript
+  import { createEnhancedFindset } from '@algorithm.ts/findset'
+
+  const findset = createEnhancedFindset(100)
+
+  findset.init(100)
+  findset.size(1)       // => 1
+  findset.merge(1, 2)
+  findset.size(1)       // => 2
+  findset.size(2)       // => 2
+  findset.getSetOf(1)   // => Set {1, 2}
+  findset.getSetOf(2)   // => Set {1, 2}
+  ```
+
 
 ## Related
 
diff --git a/packages/findset/README.md b/packages/findset/README.md
@@ -128,6 +128,25 @@ amortized constant time complexity:
   findset.size(4)   // => 1
   ```
 
+* Create an enhanced findset:
+
+  On the basis of ordinary findset, this enhanced version also supports to get
+  all the nodes on a given tree (access through the root node).
+
+  ```typescript
+  import { createEnhancedFindset } from '@algorithm.ts/findset'
+
+  const findset = createEnhancedFindset(100)
+
+  findset.init(100)
+  findset.size(1)       // => 1
+  findset.merge(1, 2)
+  findset.size(1)       // => 2
+  findset.size(2)       // => 2
+  findset.getSetOf(1)   // => Set {1, 2}
+  findset.getSetOf(2)   // => Set {1, 2}
+  ```
+
 
 ## Related
 
diff --git a/packages/findset/__test__/findset.spec.ts b/packages/findset/__test__/findset.spec.ts
@@ -1,5 +1,11 @@
+/* eslint-disable no-param-reassign */
 import { randomInt } from '@algorithm.ts/knuth-shuffle'
-import { createFindSet, createHeuristicFindSet } from '../src'
+import type { EnhancedFindset } from '../src'
+import {
+  createEnhancedFindset,
+  createFindSet,
+  createHeuristicFindSet,
+} from '../src'
 
 describe('createFindSet', function () {
   const MAX_N = 1000
@@ -91,3 +97,176 @@ describe('createHeuristicFindSet', function () {
     expect(count).toBe(MAX_N)
   })
 })
+
+describe('enhanced-findset', function () {
+  const findset = createEnhancedFindset(10)
+
+  test('init', function () {
+    findset.init(5)
+    for (let x = 1; x <= 5; ++x) {
+      expect(Array.from(findset.getSetOf(x)!)).toEqual([x])
+    }
+    for (let x = 6; x <= 10; ++x) {
+      expect(Array.from(findset.getSetOf(x)!)).toEqual([])
+    }
+
+    expect(() => findset.init(0)).toThrow(
+      /Invalid value, expect an integer in the range of/,
+    )
+    expect(() => findset.init(1)).not.toThrow()
+    expect(() => findset.init(10)).not.toThrow()
+    expect(() => findset.init(11)).toThrow(
+      /Invalid value, expect an integer in the range of/,
+    )
+  })
+
+  test('merge', function () {
+    findset.init(10)
+    findset.merge(2, 3)
+    expect(findset.size(2)).toEqual(2)
+    expect(findset.size(3)).toEqual(2)
+    expect(Array.from(findset.getSetOf(2)!).sort()).toEqual([2, 3])
+    expect(Array.from(findset.getSetOf(3)!).sort()).toEqual([2, 3])
+
+    findset.merge(1, 2)
+    expect(findset.size(1)).toEqual(3)
+    expect(findset.size(2)).toEqual(3)
+    expect(Array.from(findset.getSetOf(1)!).sort()).toEqual([1, 2, 3])
+    expect(Array.from(findset.getSetOf(2)!).sort()).toEqual([1, 2, 3])
+
+    findset.merge(1, 1)
+    expect(findset.size(1)).toEqual(3)
+    expect(Array.from(findset.getSetOf(1)!).sort()).toEqual([1, 2, 3])
+    expect(Array.from(findset.getSetOf(2)!).sort()).toEqual([1, 2, 3])
+  })
+})
+
+describe('leetcode', function () {
+  test('2092 Find All People With Secret', function () {
+    const solve = createSolve()
+    const data: Array<{ input: Parameters<typeof solve>; answer: number[] }> = [
+      {
+        input: [
+          6,
+          [
+            [1, 2, 5],
+            [2, 3, 8],
+            [1, 5, 10],
+          ],
+          1,
+        ],
+        answer: [0, 1, 2, 3, 5],
+      },
+      {
+        input: [
+          4,
+          [
+            [3, 1, 3],
+            [1, 2, 2],
+            [0, 3, 3],
+          ],
+          3,
+        ],
+        answer: [0, 1, 3],
+      },
+      {
+        input: [
+          5,
+          [
+            [3, 4, 2],
+            [1, 2, 1],
+            [2, 3, 1],
+          ],
+          1,
+        ],
+        answer: [0, 1, 2, 3, 4],
+      },
+      {
+        input: [
+          6,
+          [
+            [0, 2, 1],
+            [1, 3, 1],
+            [4, 5, 1],
+          ],
+          1,
+        ],
+        answer: [0, 1, 2, 3],
+      },
+    ]
+    for (const kase of data) {
+      const [N, meetings, firstPerson] = kase.input
+      expect(solve(N, meetings, firstPerson)).toEqual(kase.answer)
+    }
+
+    function createSolve(): (
+      N: number,
+      meetings: number[][],
+      firstPerson: number,
+    ) => number[] {
+      const MAX_N = 1e5 + 10
+      const answer: Set<number> = new Set()
+      const nodes: Set<number> = new Set()
+      const visited: Uint8Array = new Uint8Array(MAX_N)
+      const findset: EnhancedFindset = createEnhancedFindset(MAX_N)
+
+      return function findAllPeople(
+        N: number,
+        meetings: number[][],
+        firstPerson: number,
+      ): number[] {
+        const M: number = meetings.length
+
+        answer.clear()
+        answer.add(1)
+        answer.add(firstPerson + 1)
+
+        meetings
+          .sort((x, y) => x[2] - y[2])
+          .forEach(item => {
+            item[0] += 1
+            item[1] += 1
+          })
+
+        for (let i = 0, j: number; i < M; i = j) {
+          const t: number = meetings[i][2]
+          for (j = i + 1; j < M; ++j) {
+            if (meetings[j][2] !== t) break
+          }
+
+          nodes.clear()
+          for (let k = i; k < j; ++k) {
+            const [x, y] = meetings[k]
+            nodes.add(x)
+            nodes.add(y)
+          }
+
+          for (const x of nodes) {
+            findset.initNode(x)
+            visited[x] = 0
+          }
+
+          for (let k = i; k < j; ++k) {
+            const [x, y] = meetings[k]
+            findset.merge(x, y)
+          }
+
+          for (const x of nodes) {
+            if (!answer.has(x)) continue
+
+            const xx: number = findset.root(x)
+            if (visited[xx]) continue
+            visited[xx] = 1
+
+            const xxSet: Set<number> = findset.getSetOf(xx)!
+            for (const t of xxSet) answer.add(t)
+          }
+        }
+
+        return Array.from(answer)
+          .map(x => x - 1)
+          .sort((x, y) => x - y)
+      }
+    }
+  })
+})
diff --git a/packages/findset/src/enhanced.ts b/packages/findset/src/enhanced.ts
@@ -0,0 +1,95 @@
+import type { FindSet } from './ordinary'
+
+/**
+ * Enhanced findset.
+ *
+ * On the basis of ordinary findset, this enhanced version also supports to get
+ * all the nodes on a given tree (access through the root node).
+ *
+ * Example:
+ *
+ *  findset = createEnhancedFindset(20)
+ *
+ *  findset.initNode
+ *  findset.merge(2)
+ */
+export interface EnhancedFindset extends FindSet {
+  /**
+   *
+   * @param x
+   */
+  initNode(x: number): void
+  /**
+   * Size (number of the nodes) of the tree which x belongs.
+   * @param x
+   */
+  size(x: number): number
+  /**
+   *
+   * @param x
+   */
+  getSetOf(x: number): Readonly<Set<number>> | undefined
+}
+
+/**
+ * Create an enhanced findset.
+ * @param MAX_N  max nodes in the findset.
+ * @returns
+ */
+export function createEnhancedFindset(MAX_N: number): EnhancedFindset {
+  const pa: Uint32Array = new Uint32Array(MAX_N + 1)
+  const sets: Array<Set<number>> = new Array(MAX_N + 1)
+  for (let i = 0; i <= MAX_N; ++i) sets[i] = new Set()
+  return { init, root, merge, initNode, size, getSetOf }
+
+  function init(N: number): void {
+    if (N < 1 || N > MAX_N) {
+      throw new TypeError(
+        `Invalid value, expect an integer in the range of [1, ${MAX_N}], but got ${N}.`,
+      )
+    }
+
+    for (let x = 1; x <= N; ++x) initNode(x)
+  }
+
+  function root(x: number): number {
+    // eslint-disable-next-line no-return-assign
+    return pa[x] === x ? x : (pa[x] = root(pa[x]))
+  }
+
+  function merge(x: number, y: number): number {
+    let xx: number = root(x)
+    let yy: number = root(y)
+    if (xx === yy) return xx
+
+    // Heuristic combination
+    if (sets[xx].size < sets[yy].size) {
+      const t: number = xx
+      xx = yy
+      yy = t
+    }
+
+    // Merge two set.
+    for (const t of sets[yy]) sets[xx].add(t)
+    sets[yy].clear()
+
+    // eslint-disable-next-line no-return-assign
+    return (pa[yy] = xx)
+  }
+
+  function initNode(x: number): void {
+    pa[x] = x
+    sets[x].clear()
+    sets[x].add(x)
+  }
+
+  function size(x: number): number {
+    const xx: number = root(x)
+    return sets[xx].size
+  }
+
+  function getSetOf(x: number): Readonly<Set<number>> {
+    const xx: number = root(x)
+    return sets[xx]
+  }
+}
diff --git a/packages/findset/src/index.ts b/packages/findset/src/index.ts
@@ -1,2 +1,3 @@
+export * from './enhanced'
 export * from './heuristic'
 export * from './ordinary'

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
	`1`	`+export * from './enhanced'`
`1`	`2`	`export * from './heuristic'`
`2`	`3`	`export * from './ordinary'`