💥 [breaking] refactor: refactor dijkstra to make it more robust

guanghechen · guanghechen · commit 1543feac979a · 2022-01-31T13:46:35.000+08:00
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -6,6 +6,7 @@
     "gitmoji",
     "guanghechen",
     "Isap",
+    "kase",
     "leetcode",
     "lowbit",
     "manacher",
diff --git a/packages/dijkstra/README.md b/packages/dijkstra/README.md
@@ -82,53 +82,107 @@ The following definition is quoted from Wikipedia (https://en.wikipedia.org/wiki
 
 ## Usage
 
-A solution for leetcode "Number of Ways to Arrive at Destination"
-(https://leetcode.com/problems/number-of-ways-to-arrive-at-destination/):
-
-```typescript
-import type { DijkstraEdge } from '@algorithm.ts/dijkstra'
-import dijkstra from '@algorithm.ts/dijkstra'
-
-export function countPaths(N: number, roads: number[][]): number {
-  const MOD = 1e9 + 7
-  const G: Array<Array<DijkstraEdge<number>>> = new Array(N)
-  for (let i = 0; i < N; ++i) G[i] = []
-  for (const road of roads) {
-    G[road[0]].push({ to: road[1], cost: road[2] })
-    G[road[1]].push({ to: road[0], cost: road[2] })
-  }
+* Simple
+
+  ```typescript
+  import dijkstra from '@algorithm.ts/dijkstra'
+
+  const dist: number[] = dijkstra({
+    N: 4,
+    source: 0,
+    edges: [
+      { to: 1, cost: 2 },
+      { to: 2, cost: 2 },
+      { to: 3, cost: 2 },
+      { to: 3, cost: 1 },
+    ],
+    G: [[0], [1, 2], [3], []],
+  })
+  // => [0, 2, 4, 4], Which means:
+  // 
+  //      0 --> 0: cost is 0
+  //      0 --> 1: cost is 2
+  //      0 --> 2: cost is 4
+  //      0 --> 3: cost is 4
+  ```
+
+* Pass custom `dist` array.
+
+  ```typescript
+  import dijkstra from '@algorithm.ts/dijkstra'
+
+  const dist: number[] = []
+  dijkstra(
+    {
+      N: 4,
+      source: 0,
+      edges: [
+        { to: 1, cost: 2 },
+        { to: 2, cost: 2 },
+        { to: 3, cost: 2 },
+        { to: 3, cost: 1 },
+      ],
+      G: [[0], [1, 2], [3], []],
+    },
+    undefined,
+    dist
+  )
+
+  dist // => [0, 2, 4, 4]
+  ```
+
+### Example
 
-  const source = 0
-  const target = N - 1
-  const dist: number[] = dijkstra<number>(N, target, G, 0, 1e12)
+* A solution for leetcode "Number of Ways to Arrive at Destination"
+  (https://leetcode.com/problems/number-of-ways-to-arrive-at-destination/):
 
-  const dp: number[] = new Array(N).fill(-1)
-  return dfs(source)
+  ```typescript
+  import type { DijkstraEdge } from '@algorithm.ts/dijkstra'
+  import dijkstra from '@algorithm.ts/dijkstra'
+
+  export function countPaths(N: number, roads: number[][]): number {
+    const edges: DijkstraEdge[] = []
+    const G: number[][] = new Array(N)
+    for (let i = 0; i < N; ++i) G[i] = []
+    for (const [from, to, cost] of roads) {
+      G[from].push(edges.length)
+      edges.push({ to, cost })
+
+      G[to].push(edges.length)
+      edges.push({ to: from, cost })
+    }
 
-  function dfs(o: number): number {
-    if (o === target) return 1
+    const source = 0
+    const target = N - 1
+    const dist: number[] = dijkstra({ N, source: target, edges, G }, 1e12)
 
-    let answer = dp[o]
-    if (answer !== -1) return answer
+    const dp: number[] = new Array(N).fill(-1)
+    return dfs(source)
 
-    answer = 0
-    const d = dist[o]
-    for (const e of G[o]) {
-      if (dist[e.to] + e.cost === d) {
-        const t = dfs(e.to)
-        answer = modAdd(answer, t)
+    function dfs(o: number): number {
+      if (o === target) return 1
+
+      let answer = dp[o]
+      if (answer !== -1) return answer
+
+      answer = 0
+      const d = dist[o]
+      for (const idx of G[o]) {
+        const e: DijkstraEdge = edges[idx]
+        if (dist[e.to] + e.cost === d) {
+          const t = dfs(e.to)
+          answer = modAdd(answer, t)
+        }
       }
+      return dp[o] = answer
     }
-    dp[o] = answer
-    return answer
   }
 
   function modAdd(x: number, y: number): number {
     const z: number = x + y
     return z < MOD ? z : z - MOD
   }
-}
-```
+  ```
 
 
 ## Related
diff --git a/packages/dijkstra/__test__/dijkstra.spec.ts b/packages/dijkstra/__test__/dijkstra.spec.ts
@@ -1,7 +1,25 @@
-import type { DijkstraEdge } from '../src'
+import type { IEdge } from '../src'
 import dijkstra from '../src'
 
-describe('dijkstra', function () {
+describe('basic', function () {
+  test('simple', function () {
+    expect(
+      dijkstra({
+        N: 4,
+        source: 0,
+        edges: [
+          { to: 1, cost: 2 },
+          { to: 2, cost: 2 },
+          { to: 3, cost: 2 },
+          { to: 3, cost: 1 },
+        ],
+        G: [[0], [1, 2], [3], []],
+      }),
+    ).toEqual([0, 2, 4, 4])
+  })
+})
+
+describe('leetcode', function () {
   // https://leetcode.com/problems/number-of-ways-to-arrive-at-destination/
   test('leetcode/number-of-ways-to-arrive-at-destination', function () {
     const data: any[] = [
@@ -29,24 +47,32 @@ describe('dijkstra', function () {
       },
     ]
 
+    const customDist: number[] = []
     for (const kase of data) {
-      expect(countPaths(kase.input[0], kase.input[1])).toEqual(kase.answer)
+      const [N, roads] = kase.input
+
+      expect(countPaths(N, roads)).toEqual(kase.answer)
+      expect(countPaths(N, roads, customDist)).toEqual(kase.answer)
     }
   })
 })
 
 const MOD = 1e9 + 7
-function countPaths(N: number, roads: number[][]): number {
-  const G: Array<Array<DijkstraEdge<number>>> = new Array(N)
+function countPaths(N: number, roads: number[][], customDist?: number[]): number {
+  const edges: IEdge[] = []
+  const G: number[][] = new Array(N)
   for (let i = 0; i < N; ++i) G[i] = []
-  for (const road of roads) {
-    G[road[0]].push({ to: road[1], cost: road[2] })
-    G[road[1]].push({ to: road[0], cost: road[2] })
+  for (const [from, to, cost] of roads) {
+    G[from].push(edges.length)
+    edges.push({ to, cost })
+
+    G[to].push(edges.length)
+    edges.push({ to: from, cost })
   }
 
   const source = 0
   const target = N - 1
-  const dist: number[] = dijkstra<number>(N, target, G, 0, 1e12)
+  const dist: number[] = dijkstra({ N, source: target, edges, G }, 1e12, customDist)
 
   const dp: number[] = new Array(N).fill(-1)
   return dfs(source)
@@ -59,7 +85,8 @@ function countPaths(N: number, roads: number[][]): number {
 
     answer = 0
     const d = dist[o]
-    for (const e of G[o]) {
+    for (const idx of G[o]) {
+      const e: IEdge = edges[idx]
       if (dist[e.to] + e.cost === d) {
         const t = dfs(e.to)
         answer = modAdd(answer, t)
diff --git a/packages/dijkstra/src/index.ts b/packages/dijkstra/src/index.ts
@@ -1,51 +1,73 @@
 import { createPriorityQueue } from '@algorithm.ts/priority-queue'
 
-export interface DijkstraEdge<T extends number | bigint> {
+export interface IEdge {
   /**
    * The other end of the edge.
    */
   to: number
   /**
    * The cost of walking along this side.
    */
-  cost: T
+  cost: number
 }
 
+export interface IGraph {
+  /**
+   * The number of nodes in the graph. (0-index)
+   */
+  N: number
+  /**
+   * The source node. (0-index)
+   */
+  source: number
+  /**
+   * Graph edges.
+   */
+  edges: ReadonlyArray<IEdge>
+  /**
+   * Adjacency list. G[i] represent the index list of the  edges start from node i.
+   */
+  G: ReadonlyArray<ReadonlyArray<number>>
+}
+
+const ZERO = 0
+const DEFAULT_INF = Math.floor(Number.MAX_SAFE_INTEGER / 2)
+const Q = createPriorityQueue<{ pos: number; cost: number }>((x, y) => {
+  if (x.cost === y.cost) return 0
+  return x.cost < y.cost ? 1 : -1
+})
+
 /**
  * The dijkstra algorithm, optimized with priority queue.
  *
- * @param N         the number of nodes in the graph
- * @param source    the source node
- * @param G         edges of the graph. G[i] represent the edge list start from node i.
- * @param ZERO      0 for number, 0n for bigint.
- * @param INF       a big number, such as Number.MAX_SAFE_INTEGER
+ * @param INF         A big number, representing the unreachable cost.
+ * @param customDist
  * @returns
  *
  * @see https://me.guanghechen.com/post/algorithm/graph/shortest-path/dijkstra
  */
-export function dijkstra<T extends number | bigint>(
-  N: number,
-  source: number,
-  G: ReadonlyArray<ReadonlyArray<DijkstraEdge<T>>>,
-  ZERO: T,
-  INF: T,
-): T[] {
-  const dist: T[] = new Array(N).fill(INF)
-  const Q = createPriorityQueue<{ pos: number; cost: T }>((x, y) => y.cost - x.cost)
+export function dijkstra(
+  graph: IGraph,
+  INF: number = DEFAULT_INF,
+  customDist?: number[],
+): number[] {
+  const { N, source, edges, G } = graph
+  const dist: number[] = customDist ?? []
+  if (dist.length < N) dist.length = N
 
-  // eslint-disable-next-line no-param-reassign
+  dist.fill(INF, 0, N)
   dist[source] = ZERO
   Q.enqueue({ pos: source, cost: ZERO })
 
   while (Q.size() > 0) {
     const { pos, cost } = Q.dequeue()!
     if (dist[pos] < cost) continue
-    for (const e of G[pos]) {
-      const candidate: T = (dist[pos] as any) + e.cost
-      if (dist[e.to] > candidate) {
-        // eslint-disable-next-line no-param-reassign
-        dist[e.to] = candidate
-        Q.enqueue({ pos: e.to, cost: dist[e.to] })
+    for (const idx of G[pos]) {
+      const edge: IEdge = edges[idx]
+      const candidate: number = dist[pos] + edge.cost
+      if (dist[edge.to] > candidate) {
+        dist[edge.to] = candidate
+        Q.enqueue({ pos: edge.to, cost: dist[edge.to] })
       }
     }
   }
