描述
现在有红,绿两种颜色的石头,现在我们需要用这两种石头搭建一个塔,塔需要满足如下三个条件:
- 第1层应该包含1块石头,第2层应该包含2块,第i层需要包含i块石头。
- 同一层的石头应该是同一个颜色(红或绿)。
- 塔的层数尽可能多。
在满足上面三个条件的前提下,有多少种不同的建造塔的方案?当塔中任意一个对应位置的石头颜色不同,我们就认为这两个方案不相同。石头可以不用完。
由于答案可能会很大,请对10^9+7取模。
- red+green≥1
- 0≤red,green≤6×104
在线评测地址:领扣题库官网
样例1
输入: 4 6
输出: 2
说明: 有两种方案:[红,绿,红,绿],[绿,绿,绿,红]
源代码
public int twoColorsTower(int red, int green) {
if (red == 0 || green == 0) {
return 1;
}
if (red > green) {
int temp = red;
red = green;
green = temp;
}
// dp[i&1][j]表示前i层一共放j个红石头
int[][] dp = new int[2][red + 1];
dp[1][0] = dp[1][1] = 1;
int level = (int) Math.sqrt(2 * (red + green));
int sum = 1, lower = 0, upper = red;
int curr = 1, prev;
int MOD = (int) 1.0e9 + 7;
for (int i = 2; i <= level; i++) {
sum += i;
int tmpUpper = Math.min(sum, red);
int tmpLower = Math.max(sum - green, 0);
// 红石头不够了,已经是最高层,停止更新
if (tmpLower > tmpUpper) break;
upper = tmpUpper;
lower = tmpLower;
prev = curr;
curr = curr ^ 1;
// j小于本层i,红石只能是之前都放完了
for (int j = lower; j < i; j++) {
dp[curr][j] = dp[prev][j];
}
// 转移方程:dp[i][j] = dp[i-1][j] + dp[i-1][j-i] (when j>=i)
// dp[i-1][j]表示第i层放绿石 dp[i-1][j-i]表示i层放红石
for (int j = i; j <= upper; j++) {
dp[curr][j] = (dp[prev][j] + dp[prev][j - i]) % MOD;
}
}
int ans = 0;
for (int j = lower; j <= upper; j++) {
ans = (ans + dp[curr][j]) % MOD;
}
return ans;
}
更多题解参考:九章官网solution