数组转树和树转数组

前端经常遇到树的扁平化以及把扁平化的数组转成树。本文各实现了一种不用递归的方法。其中array2Tree()方法当只有一个根节点时最大时间复杂度O(n^2)。tree2Array()方法当只有一个根节点时时间复杂度O(n)。

树结构可以是有多个根节点的情况，比如

type Item = {
  id: string | number
  pid: string | number
  children?: Item[]
  [k: string]: unknown
}
export function array2Tree(array: Item[]) {
  // 收集所有节点的id
  const allIdSet = new Set() 
  for (let i = 0, len = array.length; i < len; i++) {
    allIdSet.add(array[i].id)
  }
  // 收集所有的根节点，有可能不止一个根节点
  const rootNodes = []
  for (let i = 0, len = array.length; i < len; i++) {
    if (!allIdSet.has(array[i].pid)) {
      rootNodes.push(array[i])
    }
  }
  // 挨个处理每个根节点
  for (let k = 0, kLen = rootNodes.length; k < kLen; k++) {
    const queue = [rootNodes[k]]
    // 借助queue队列按广度优先遍历的顺序一层一层收集每个节点的children
    while (queue.length) {
      const p = queue.shift() 
      const children = []
      // 找p的所有子节点
      for (let i = 0, len = array.length; i < len; i++) {
        if (array[i].pid === p.id) {
          children.push(array[i])
          queue.push(array[i])
        }
      }
      p.children = children
    }
  }

  return rootNodes
}

/**
 * tree 数组中可能有多个根节点
 * @param tree 
 */
export function tree2Array(tree: Item[]) {
  const result = []
  // 这层循环挨个处理每个根节点
  for (let i = 0, len = tree.length; i < len; i++) {
    const cur = tree[i]
    const children = cur.children
    // 没有children的孤立根节点
    if (!children?.length) {
      result.push({
        id: cur.id,
        pid: cur.pid,
      })
      continue
    }

    // 借助queue队列按深度优先遍历的顺序依次让children[i]这条链上的节点进队列出队列，
    // 再同样的处理children[i+1]这条链。
    const queue = [{node: cur, index: 0}]
    // 标记 queue[queue.length - 1] 所取到的值是正向新加进去的，还是由于把子孙节点处理完了 后退到这的
    let newAdd = true 
    while (queue.length) {
      const cur = queue[queue.length - 1]
      const children = cur.node.children
      let index = 0
      // 表示queue队列中的cur节点的children[index]这条子节点链上的都排出队列了，又退回到cur节点，
      // 接下来该处理cur的children[index + 1]这条子节点链
      if (!newAdd) {
        index = cur.index + 1
        cur.index = index
      }
      if (children?.length && index < children.length) {
        newAdd = true
        queue.push({
          node: children[index],
          index: index,
        })
      } else {
        newAdd = false
        const t = queue.pop()
        result.push({
          id: t.node.id,
          pid: t.node.pid,
        })
      }
    }
  }

  return result
}

// test case
const array = [
  {id: 1, name: '部门1', pid: 0},
  {id: 2, name: '部门2', pid: 1},
  {id: 3, name: '部门3', pid: 1},
  {id: 4, name: '部门3', pid: 0},
  {id: 5, name: '部门4', pid: 4},
  {id: 6, name: '部门5', pid: 0},
  {id: 7, name: '部门5', pid: 6},
  {id: 8, name: '部门5', pid: 7},
  {id: 9, name: '部门9', pid: 0},
]

const result = array2Tree(array) 
console.log(JSON.stringify(result))
const result2 = tree2Array(result)
console.log(JSON.stringify(result2))