feat(TreeData): 遍历算法优化

fjc0k · Dec 14, 2020 · e7a8cf8 · e7a8cf8
1 parent e238dbf
commit e7a8cf8
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 55 deletions.
diff --git a/src/utils/TreeData.test.ts b/src/utils/TreeData.test.ts
@@ -203,6 +203,23 @@ describe('TreeData', () => {
     ).toMatchSnapshot()
   })
 
+  // test('大数据', () => {
+  //   const data = range(0, 1000).map(i => ({
+  //     id: `${i}`,
+  //     children: range(0, 1000).map(i2 => ({
+  //       id: `${i}.${i2}`,
+  //       children: [],
+  //     })),
+  //   }))
+  //   const startTime = Date.now()
+  //   new TreeData(data).traverse(_ => {
+  //     // ...
+  //   })
+  //   const endTime = Date.now()
+  //   console.log(endTime - startTime)
+  //   expect(1).toBe(1)
+  // })
+
   test('综合', () => {
     const names: string[] = []
     const names2: string[] = []

diff --git a/src/utils/TreeData.ts b/src/utils/TreeData.ts
@@ -113,17 +113,24 @@ export class TreeData<TNode extends TreeDataNode> {
   private static traverse<TNode extends TreeDataNode>(
     data: TreeDataData<TNode>,
     childrenPropName: TreeDataChildrenPropName<TNode>,
-    parentNode: TNode | undefined,
-    fn: TreeDataTraverseFn<TNode>,
     searchStrategy: TreeDataSearchStrategy,
-    depth: number,
-    path: TNode[],
+    fn: TreeDataTraverseFn<TNode>,
   ) {
-    const removeNodeIndexes: number[] = []
-    const skipChildrenTraverseNodeIndexes: number[] = []
-
-    for (let i = 0, len = data.length; i < len; i++) {
-      const node = data[i]
+    const nodes: Array<
+      [
+        node: TNode,
+        index: number,
+        parentNode: TNode | undefined,
+        siblings: TNode[],
+        depth: number,
+        path: TNode[],
+      ]
+    > = data.map((child, index) => [child, index, undefined, data, 0, []])
+
+    let currentNode: typeof nodes[0] | undefined
+    const removeNodes: Array<[siblings: TNode[], indexes: number[]]> = []
+    while ((currentNode = nodes.shift())) {
+      const [node, index, parentNode, siblings, depth, path] = currentNode
 
       let isRemove = false
       let isExit = false
@@ -136,59 +143,51 @@ export class TreeData<TNode extends TreeDataNode> {
         path: path,
         removeNode: () => {
           isRemove = true
-          removeNodeIndexes.push(i)
         },
-        exit: () => (isExit = true),
+        exit: () => {
+          isExit = true
+        },
         skipChildrenTraverse: () => {
           isSkipChildrenTraverse = true
-          skipChildrenTraverseNodeIndexes.push(i)
         },
       })
 
-      if (isExit) return
-
-      if (!isSkipChildrenTraverse) {
-        if (searchStrategy === 'DFS') {
-          if (
-            !isRemove &&
-            node[childrenPropName] &&
-            Array.isArray(node[childrenPropName])
-          ) {
-            TreeData.traverse(
-              node[childrenPropName],
-              childrenPropName,
-              node,
-              fn,
-              searchStrategy,
-              depth + 1,
-              path.concat(node),
-            )
-          }
+      if (isRemove) {
+        if (
+          !removeNodes.length ||
+          removeNodes[removeNodes.length - 1][0] !== siblings
+        ) {
+          removeNodes.push([siblings, []])
         }
+        removeNodes[removeNodes.length - 1][1].push(index)
       }
-    }
 
-    // 删除节点
-    for (let i = removeNodeIndexes.length - 1; i >= 0; i--) {
-      data.splice(removeNodeIndexes[i], 1)
+      if (isExit) return
+
+      if (
+        !isRemove &&
+        !isSkipChildrenTraverse &&
+        node[childrenPropName] &&
+        Array.isArray(node[childrenPropName])
+      ) {
+        nodes[searchStrategy === 'DFS' ? 'unshift' : 'push'](
+          ...node[childrenPropName].map((child: any, index: number) => [
+            child,
+            index,
+            node,
+            node[childrenPropName],
+            depth + 1,
+            path.concat(node),
+          ]),
+        )
+      }
     }
 
-    if (searchStrategy === 'BFS') {
-      for (let i = 0, len = data.length; i < len; i++) {
-        if (skipChildrenTraverseNodeIndexes.indexOf(i) === -1) {
-          const node = data[i]
-          if (node[childrenPropName] && Array.isArray(node[childrenPropName])) {
-            TreeData.traverse(
-              node[childrenPropName],
-              childrenPropName,
-              node,
-              fn,
-              searchStrategy,
-              depth + 1,
-              path.concat(node),
-            )
-          }
-        }
+    let removeNode: typeof removeNodes[0] | undefined
+    while ((removeNode = removeNodes.shift())) {
+      let removeNodeIndex: number | undefined
+      while ((removeNodeIndex = removeNode[1].pop()) != null) {
+        removeNode[0].splice(removeNodeIndex, 1)
       }
     }
   }
@@ -210,11 +209,8 @@ export class TreeData<TNode extends TreeDataNode> {
       TreeData.traverse<TNode>(
         this.data,
         this.childrenPropName,
-        undefined,
-        fn,
         searchStrategy,
-        0,
-        [],
+        fn,
       )
     }
     return this