fix(findNodes): shouldn't judge answer length in dfs

opposition: new Array(10).length === 10 // true
lbwa · Oct 19, 2020 · 4d01077 · 4d01077
1 parent 42358e6
commit 4d01077
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 24 deletions.
diff --git a/src/find.spec.ts b/src/find.spec.ts
@@ -11,28 +11,68 @@ describe('N-ary tree', () => {
       value: 1,
       children: [
         {
-          value: 11,
+          value: 1,
           children: [
             {
-              value: 111
+              value: 11
+            },
+            {
+              value: 12,
+              children: [
+                {
+                  value: 121
+                },
+                {
+                  value: 122
+                }
+              ]
+            },
+            {
+              value: 13
+            },
+            {
+              value: 14,
+              children: [
+                {
+                  value: 141
+                },
+                {
+                  value: 142
+                }
+              ]
             }
           ]
         },
         {
-          value: 12,
-          children: [
-            {
-              value: 121
-            }
-          ]
+          value: 2
         }
       ]
     }
-    expect(findNodes(tree, [111])).toEqual({
-      nodes: [{ value: 111 }],
-      values: [111]
+    expect(findNodes(undefined as any, [11])).toEqual({ nodes: [], values: [] })
+    expect(findNodes(tree, [11])).toEqual({
+      nodes: [{ value: 11 }],
+      values: [11]
+    })
+    expect(findNodes(tree, [142])).toEqual({
+      nodes: [{ value: 142 }],
+      values: [142]
     })
     expect(findNodes(tree, [123])).toEqual({ nodes: [], values: [] })
+    expect(
+      findNodes(
+        {
+          value: 1,
+          children: new Array(99).fill(undefined).concat({
+            value: 12,
+            children: new Array(999)
+              .fill(undefined)
+              .concat({ value: 19 })
+              .concat(new Array(99).fill(null))
+          })
+        },
+        [19]
+      )
+    ).toEqual({ nodes: [{ value: 19 }], values: [19] })
   })
 
   it('should find nodes with specific fields', () => {

diff --git a/src/find.ts b/src/find.ts
@@ -28,21 +28,14 @@ export function findNodes<N extends TreeNode<V>, V = any>(
       return
     }
 
-    let i = targets.length
-    while (i--) {
-      if (targets[i] === root[value]) {
-        values[i] = root[value]
-        nodes[i] = root
-      }
+    const index = targets.indexOf(root[value])
+    if (index > -1) {
+      values[index] = root[value]
+      nodes[index] = root
     }
 
-    // avoid unnecessary comparison, return answer as soon as possible
-    if (values.length === targets.length) {
-      return
-    }
-
-    if (Array.isArray(root[children]) && (root[children] as any).length > 0) {
-      ;(root[children] as any).forEach(dfs)
+    if (Array.isArray(root[children]) && root[children].length > 0) {
+      root[children].forEach(dfs)
     }
   }
   dfs(root)