带因子的二叉树

823. 带因子的二叉树 - 力扣（LeetCode）

给出一个含有不重复整数元素的数组 arr ，每个整数 arr[i] 均大于 1。

用这些整数来构建二叉树，每个整数可以使用任意次数。其中：每个非叶结点的值应等于它的两个子结点的值的乘积。

满足条件的二叉树一共有多少个？答案可能很大，返回 对 10^9 + 7 取余 的结果。

示例 1:

输入: arr = [2, 4]
输出: 3
解释: 可以得到这些二叉树: [2], [4], [4, 2, 2]

示例 2:

输入: arr = [2, 4, 5, 10]
输出: 7
解释: 可以得到这些二叉树: [2], [4], [5], [10], [4, 2, 2], [10, 2, 5], [10, 5, 2].

提示：

• 1 <= arr.length <= 1000
• 2 <= arr[i] <= 10^9
• arr 中的所有值 互不相同

思路

本题我们可以用动态规划求解。首先我们对 arr 进行升序排序，排序后以 arr[i]为根节点的构造的带因子的二叉树，他的子节点必定小于 arr[i],在它的左侧。我们用 dp[i]来表示以 arr[i]为根节点的二叉树的个数，假设 arr[j] \* arr[k] = arr[i],则我们可以用 arr[j] 和 arr[k] 分别为左右子节点来构造二叉树。

• 如果 j === k, 一共有 dp[j] * dp[k] 个二叉树
• 如果 j !== k,左右节点可以互换，一共有 dp[j] * dp[k] * 2 个二叉树

求 dp[i] 时，只需要用双指针从两端遍历 arr[0] 到 arr[i-1],找到符合要求的二叉树，并把他们求和，因为 arr[i] 本身也是符合要求的二叉树，最后 dp[i] += 1。初始状态 dp[0] = 1。求出所有 dp 后，再对他们求和，就是我们要求的解。

解题

/**
 * @param {number[]} arr
 * @return {number}
 */
var numFactoredBinaryTrees = function (arr) {
  const MOD = 10 ** 9 + 7;

  const n = arr.length;
  const dp = Array(n).fill(1);
  arr.sort((a, b) => a - b);
  for (let i = 1; i < n; i++) {
    let l = 0,
      r = i - 1;
    while (l <= r) {
      let m = arr[l] * arr[r];
      if (m === arr[i]) {
        let k = (dp[l] * dp[r] * (l !== r ? 2 : 1)) % MOD;
        dp[i] = (dp[i] + k) % MOD;
        l++;
        r--;
      } else if (m > arr[i]) {
        r--;
      } else {
        l++;
      }
    }
  }
  return dp.reduce((sum, num) => sum + num) % MOD;
};