Sum of K Subarrays With Length at Least M - Problem

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>

#define MAX(a,b) (((a)>(b))?(a):(b))
#define MIN(a,b) (((a)<(b))?(a):(b))

long long solution(int* nums, int n, int k, int m) {
    if (n < k * m) return 0;

    // Build prefix sum
    long long* prefix = (long long*)calloc(n + 1, sizeof(long long));
    for (int i = 0; i < n; i++) {
        prefix[i + 1] = prefix[i] + nums[i];
    }

    // dp[i][j] = max sum using j subarrays from first i elements
    // Using a large negative number slightly safe from underflow when adding
    long long INF = LLONG_MIN / 2;
    
    // Allocate DP table (n+1) x (k+1)
    long long** dp = (long long**)malloc((n + 1) * sizeof(long long*));
    for (int i = 0; i <= n; i++) {
        dp[i] = (long long*)malloc((k + 1) * sizeof(long long));
        for (int j = 0; j <= k; j++) {
            dp[i][j] = INF;
        }
        dp[i][0] = 0;
    }

    for (int i = m; i <= n; i++) {
        // We can't have more than i/m segments
        int max_j = i / m;
        if (max_j > k) max_j = k;

        for (int j = 1; j <= max_j; j++) {
            // Option 1: don't include element i-1 in the j-th subarray end
            dp[i][j] = dp[i-1][j];

            // Option 2: end a subarray at position i-1 (1-based index i)
            // The loop condition length <= i - j + 1 ensures we leave space for j-1 previous items
            // (Strictly speaking for size m, it relies on the dp value check of previous state)
            for (int length = m; length <= i - j + 1; length++) {
                if (i - length >= 0) {
                    long long subarraySum = prefix[i] - prefix[i - length];
                    long long prev = dp[i - length][j - 1];
                    
                    if (prev > INF) { // Only add if previous state is valid
                        dp[i][j] = MAX(dp[i][j], prev + subarraySum);
                    }
                }
            }
        }
    }

    long long result = dp[n][k];
    
    // Cleanup
    free(prefix);
    for (int i = 0; i <= n; i++) free(dp[i]);
    free(dp);

    return result == INF ? 0 : result;
}

int main() {
    char line[10000];
    
    // Read nums string
    if (!fgets(line, sizeof(line), stdin)) return 0;
    
    // Parse [1,2,3]
    int* nums = (int*)malloc(sizeof(int) * 5000);
    int count = 0;
    char* ptr = line;
    while (*ptr) {
        if (*ptr == '[' || *ptr == ']' || *ptr == ',') *ptr = ' ';
        ptr++;
    }
    
    ptr = line;
    char* end;
    while (1) {
        long val = strtol(ptr, &end, 10);
        if (ptr == end) break;
        nums[count++] = (int)val;
        ptr = end;
    }

    // Read k
    if (!fgets(line, sizeof(line), stdin)) return 0;
    int k = atoi(line);

    // Read m
    if (!fgets(line, sizeof(line), stdin)) return 0;
    int m = atoi(line);

    printf("%lld\n", solution(nums, count, k, m));

    free(nums);
    return 0;
}