题解 | 三角形取数(Hard Version)

/*
 思路：先处理三角形数塔的输入，然后把数塔赋值给dp数组。dp数组要先处理最下面一层，把|j - center| > k的赋值为INF。然后两层for循环往上推，最后dp【0】【0】就是结果。
 细节：注意for循环的边界问题
*/


#include <iostream>
#include<vector>
#include<algorithm>
#include<cmath>
using namespace std;

const long long INF = -1e18;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n,k;
    cin>>n>>k;
    vector<vector<long long>>arr(n);
    for(int i = 0; i <= n -1; i++){
        arr[i].resize(2 * i + 1);
        for(int j = 0; j < 2 * i + 1; j++){
            cin>>arr[i][j];
        }
    }
    vector<vector<long long>>dp = arr;
    int center = n -1;
    for(int j = 0; j < 2 * (n - 1) + 1; j ++){
        if(abs(j - center) > k){
            dp[n -1][j] = INF;
        }
    }
    for(int i = n -2; i >= 0; i--){
        for(int j = 0; j < 2 * i + 1; j++){
            long long next_max = max({dp[i + 1][j],dp[i + 1][j + 1],dp[i + 1][j + 2]});
            if(next_max != INF){
                dp[i][j] += next_max;
            }
            else dp[i][j] = INF;
        }
    }
    long long res = dp[0][0];
    cout<<res<<'\n';

}
// 64 位输出请用 printf("%lld")