TypeScript implementation of modified KMP algorithm for tree model

matchTree.ts

export type Tree<T, K> = T & {
  id: K
  children: Tree<T, K>[]
}

// Based on KMP algorithm
export function matchTree<T, K>(
  root: Tree<T>,
  prop: (value: Tree<T>) => K,
  selector: K[],
) {
  const lps = createLPS(selector)
  const st = [[root, 0]] as [Tree<T>, number][]
  const results = [] as Tree<T>[]

  while (st.length) {
    const [node, j] = st.pop()!

    if (prop(node) === selector[j]) {
      if (j === selector.length - 1) {
        results.push(node)
        for (const child of reverse(node.children)) {
          st.push([child, lps[j]])
        }
        continue
      }

      for (const child of reverse(node.children)) {
        st.push([child, j + 1])
      }
      continue
    }

    if (j === 0) {
      for (const child of reverse(node.children)) {
        st.push([child, j])
      }
    } else {
      st.push([node, lps[j - 1]])
    }
  }

  return results
}

function createLPS<K>(selector: K[]) {
  const lps: number[] = new Array(selector.length).fill(0)
  let length = 0
  let i = 1

  while (i < selector.length) {
    if (selector[i] === selector[length]) {
      length++
      lps[i] = length
      i++
    } else {
      if (length !== 0) {
        length = lps[length - 1]
      } else {
        lps[i] = 0
        i++
      }
    }
  }
  return lps
}