1483.树节点的第-k-个祖先.md

1483.树节点的第-k-个祖先

/*
 * @lc app=leetcode.cn id=1483 lang=typescript
 *
 * [1483] 树节点的第 K 个祖先
 */

// @lc code=start
class TreeAncestor {
  constructor(n: number, p: number[]) {}

  getKthAncestor(node: number, k: number): number {}
}

/**
 * Your TreeAncestor object will be instantiated and called as such:
 * var obj = new TreeAncestor(n, parent)
 * var param_1 = obj.getKthAncestor(node,k)
 */
// @lc code=end

解法 1: 树上倍增

class TreeAncestor {
  constructor(n: number, p: number[]) {
    const st = Array.from({ length: n }, () => [])

    const h = new Array(n).fill(-1),
      e: number[] = [],
      ne: number[] = []
    const add = (i: number, j: number) => (e.push(j), ne.push(h[i]), (h[i] = ne.length - 1))
    for (let i = 1; i < n; i++) add(p[i], i)
    const dfs = (u: number, p: number) => {
      if (p !== -1) st[u][0] = p
      for (let i = 1; st[st[u][i - 1]]?.[i - 1] !== undefined; i++) {
        st[u][i] = st[st[u][i - 1]][i - 1]
      }
      for (let i = h[u]; ~i; i = ne[i]) {
        const v = e[i]
        dfs(v, u)
      }
    }
    dfs(0, -1)
    this.st = st
  }

  getKthAncestor(node: number, k: number): number {
    for (let i = 0; i < 20; i++) {
      if (k & (1 << i)) {
        node = this.st[node][i]
        if (node === undefined) return -1
      }
    }
    return node
  }
}

Case

test.each([
  {
    input: {
      op: ['TreeAncestor', 'getKthAncestor', 'getKthAncestor', 'getKthAncestor'],
      param: [
        [9, [-1, 0, 0, 1, 1, 2, 2, 6, 7]],
        [3, 1],
        [5, 2],
        [6, 3],
      ],
    },
    output: [null, 1, 0, -1],
  },
])('input: param ', ({ input: { op, param }, output }) => {
  const treeAncestor = new TreeAncestor(...param[0])
  let res = [null]
  for (let i = 1; i < op.length; i++) {
    const key = op[i]
    res.push(treeAncestor[key](...param[i]))
  }
  expect(res).toEqual(output)
})