USACO奶牛谱系问题：C++唯一二叉树计数代码错误排查

问题描述 #

统计节点数为N、高度为K的唯一二叉树数量（结果模9901）。现有代码在小输入（如N=5,K=3、N=9,K=4）下输出正确，但N=35,K=7时输出7642，正确结果应为5024。

原代码 #

#include <bits/stdc++.h>

using namespace std;

int main() {
  // values
  int N, K;
  int possibilities[201][101]{};

  // input
  ifstream fin("nocows.in");
  fin >> N >> K;

  // solution
  possibilities[1][1] = 1;
  for (int k = 2; k <= K; k++){ // for every possible height
    // 1 node is used for the "mother", the amount of the  rest of the nodes will be between 2*(k - 1) and -2 * (1-2^(k-1))
    for (int n = 2 * (k - 1); n <= 2 * (pow(2, k - 1) - 1) && n <= N; n += 2 ){
      // one side will have n1 nodes - anywhere odd number from 1 to n; and a height of k - 1
      for (int n1 = 1; n1 < n; n1 += 2){
        // the other side will have any height k - 1 or less and the remaining nodes
        for (int k1 = 1; k1 < k; k1++){
          // if they are different numbers, they will be doubled because they are not interchangable so can work for both right and left
          if (n1 != n - n1 || k1 != k - 1) 
            possibilities[n + 1][k] += (2 * possibilities[n1][k - 1] * possibilities[n - n1][k1]) % 9901;
          else
            possibilities[n + 1][k] += (possibilities[n1][k - 1] * possibilities[n - n1][k1]) % 9901;
          possibilities[n + 1][k] %= 9901;
        }
      }
    }
  }

  // output
  ofstream fout("nocows.out");
  fout << possibilities[N][K] % 9901 << '\n';

  return 0;
}

错误分析 #

1. 重复计数问题：原代码循环枚举左子树节点数n1从1到n-1（步长2），当n1和n-n1为不同值时，会分别计算左n1右n-n1、左n-n1右n1两种情况，且每次都乘以2，导致同一组合被重复计算两次，结果错误翻倍。
2. 条件判断逻辑错误：if (n1 != n - n1 || k1 != k - 1)的判断无法避免重复累加，当右子树高度为k-1时，左右子树均为k-1的情况会被多次计入。
3. 未利用前缀和简化计算：直接枚举子树高度增加了复杂度且易出错，更高效的方式是先计算高度≤k的树的总数，再通过差值得到恰好高度为k的数量。

修正方案 #

采用前缀和递推的方式，定义两个数组：

• total[n][k]：节点数为n、高度≤k的二叉树数量
• dp[n][k]：节点数为n、高度恰好为k的二叉树数量

递推关系：

1. total[n][k]：根节点占1，左右子树节点数和为n-1（均为奇数），每个子树高度≤k-1，因此：
   total[n][k] = sum_{a=1, a为奇数}^{n-1} total[a][k-1] * total[n-1-a][k-1]，结果模9901，同时继承total[n][k-1]的结果（高度≤k包含高度≤k-1的情况）。
2. dp[n][k]：恰好高度为k的数量 = 高度≤k的总数 - 高度≤k-1的总数，为避免负数需加9901再取模：
   dp[n][k] = (total[n][k] - total[n][k-1] + 9901) % 9901

初始状态：

• total[1][k] = 1（任意k≥1，单节点树高度≤k）
• 当n < 2k-1时，dp[n][k] = 0（节点数不足，无法达到高度k）

修正后的代码 #

#include <bits/stdc++.h>
using namespace std;

const int MOD = 9901;
int total[201][101] = {0}; // total[n][k]: 节点数n，高度<=k的树的数量
int dp[201][101] = {0};    // dp[n][k]: 节点数n，高度恰好为k的树的数量

int main() {
    int N, K;
    ifstream fin("nocows.in");
    fin >> N >> K;

    // 初始化：单节点树，任意高度<=k都有1种
    for (int k = 1; k <= K; ++k) {
        total[1][k] = 1;
    }

    // 计算total数组
    for (int k = 2; k <= K; ++k) {
        // 先复制上一层的total，保证高度<=k包含高度<=k-1的情况
        for (int n = 1; n <= N; ++n) {
            total[n][k] = total[n][k-1];
        }
        // 新增高度恰好为k的情况到total中
        for (int n = 2*k -1; n <= N; n += 2) { // 节点数至少为2k-1才能达到高度k
            for (int a = 1; a < n-1; a += 2) { // 左子树节点数a，右子树节点数n-1-a
                total[n][k] = (total[n][k] + 1LL * total[a][k-1] * total[n-1 - a][k-1]) % MOD;
            }
        }
    }

    // 计算dp数组
    for (int k = 1; k <= K; ++k) {
        for (int n = 1; n <= N; ++n) {
            if (k == 1) {
                dp[n][k] = (n == 1) ? 1 : 0;
            } else {
                dp[n][k] = (total[n][k] - total[n][k-1] + MOD) % MOD;
            }
        }
    }

    ofstream fout("nocows.out");
    fout << dp[N][K] << endl;
    return 0;
}

代码说明 #

• 使用1LL强制类型转换避免整数溢出，防止两个int相乘超出范围。
• 先复制上一层total再新增当前高度的情况，逻辑更清晰，确保total[n][k]正确包含所有高度≤k的树。
• 严格限制节点数范围为2k-1到N（步长2），避免无效计算。

内容的提问来源于stack exchange，提问作者msush