2024-11-17 13:26:33

数组转树

build-tree.ts
          type DeepTreeNode<T> = T & { children: Array<DeepTreeNode<T>> }

interface BuildTreeOptions<T> {
  key: keyof T
  parentKey: keyof T
}

function buildTree<T extends object>(array: T[], options: BuildTreeOptions<T>) {
  const { key, parentKey } = options
  const tree: Array<DeepTreeNode<T>> = []
  const lookup = new Map<T[keyof T], DeepTreeNode<T>>()

  for (const item of array) {
    lookup.set(item[key], { ...item, children: [] })
  }

  for (const item of array) {
    if (item[parentKey]) {
      const parent = lookup.get(item[parentKey])
      if (parent) {
        parent.children.push(lookup.get(item[key]) as DeepTreeNode<T>)
      }
    } else {
      tree.push(lookup.get(item[key]) as DeepTreeNode<T>)
    }
  }

  return tree
}

        

浅层树(只保留一层子节点)

build-shallow-tree.ts
          type TreeNode<T, C = any> = T & { children: C[] }
type ShallowTreeNode<T> = TreeNode<T, T>

function buildShallowTree<T extends object>(array: T[], options: BuildTreeOptions<T>) {
  const { key, parentKey } = options

  const lookup = new Map<keyof T, ShallowTreeNode<T>>()

  for (const item of array) {
    lookup.set(item[key] as keyof T, { ...item, children: [] })
  }

  function* ancestry(id: keyof T): Generator<keyof T> {
    if (id && lookup.has(id)) {
      yield id
      yield* ancestry(lookup.get(id)![parentKey] as keyof T)
    }
  }

  array.forEach(node => {
    const [ancestor] = [...ancestry(node[parentKey] as keyof T)].reverse()
    if (ancestor && lookup.has(ancestor)) {
      lookup.get(ancestor)!.children.push(lookup.get(node[key] as keyof T)!)
    }
  })

  return Array.from(lookup.values()).filter(node => node[parentKey] === null)
}

        

统计树节点数量

count-tree-nodes.ts
          function countTreeNodes<T>(nodes: Array<TreeNode<T>>): number {
  let count = 0

  const traverse = (node: TreeNode<T>) => {
    count++
    for (const child of node.children) {
      traverse(child)
    }
  }

  for (const node of nodes) {
    traverse(node)
  }

  return count
}

        

树扁平化

flatten-tree.ts
          function flattenTree<T extends { children?: T[] }>(tree?: T[]): T[] {
  const result: T[] = []

  function recurse(nodes: T[] | undefined) {
    if (!nodes) return

    for (const node of nodes) {
      result.push(node)
      recurse(node.children)
    }
  }

  recurse(tree)
  return result
}