We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
  1. Prepare
  2. Algorithms
  3. Dynamic Programming
  4. Construct the Array
  5. Discussions

Construct the Array

  • cyrusxyl
     + 0 comments

    Stage 1: Brute-Force Dynamic Programming (O(NK) Time, O(NK) Space)

    Intuition: We want to find the number of valid arrays of length n ending with x. A natural approach is to build up the solution from smaller arrays. We can define dp[i][j] as the number of ways to form an array of length i where the last element is j. To form an array of length i ending in j, the previous element (at length i-1) must be different from j. Thus, dp[i][j] is the sum of dp[i-1][p] for all p from 1 to k where p != j.

    Code:

    long long countArray(int n, int k, int x) {
    long long MOD = 1000000007;

    // dp[i][j] = number of arrays of length i+1 ending with j
    // (Note: i is 0-indexed for length here, so dp[i] is for length i+1)
    std::vector<std::vector<long long>> dp(n, std::vector<long long>(k + 1, 0));

    // Base case: Array of length 1 (i=0)
    // The first element is fixed to 1.
    dp[0][1] = 1;

    for (int i = 1; i < n; ++i) { // Iterating through lengths (from 2 to n)
        long long sum_prev_row = 0;
        // Calculate sum of ways for previous length (length i)
        for (int prev_val = 1; prev_val <= k; ++prev_val) {
            sum_prev_row = (sum_prev_row + dp[i-1][prev_val]) % MOD;
        }

        for (int current_val = 1; current_val <= k; ++current_val) {
            // Number of ways to end with current_val is total ways from previous length
            // MINUS ways where previous_val was also current_val.
            dp[i][current_val] = (sum_prev_row - dp[i-1][current_val] + MOD) % MOD;
        }
    }

    return dp[n-1][x];
}

    Minor Optimizations Instead of looping to sum previous row, we can keep a separate sum variable, populate it directly when we were calculating the previous row.

    Stage 2: Space Optimization (O(NK) Time, O(K) Space)

    Intuition: Since the current row dp[i] only depends on the previous row dp[i-1], we can optimize the space by only storing the counts for the previous length. We'll use two vectors, prev and curr, where prev holds the counts for length L and curr computes counts for length L+1. After computing curr, it becomes the new prev for the next iteration.

    Code:

    long long countArray(int n, int k, int x) {
    long long MOD = 1000000007;

    // prev[j] stores ways to form array of length 'current_length - 1' ending with 'j'
    std::vector<long long> prev(k + 1, 0);
    
    // Base case: Array of length 1
    prev[1] = 1; 

    // prev_sum holds the total number of ways for the previous length
    long long prev_sum = 1; // Initially, only prev[1] is 1

    for (int i = 1; i < n; ++i) { // i represents current length - 1 (looping for lengths 2 to n)
        std::vector<long long> curr(k + 1, 0); // For the current length (i+1)
        long long current_sum = 0; 

        for (int a = 1; a <= k; ++a) {
            // To form an array of length (i+1) ending with 'a':
            // Sum of all ways for length 'i' (prev_sum)
            // Minus ways where the previous element was also 'a' (prev[a])
            curr[a] = (prev_sum - prev[a] + MOD) % MOD; 
            current_sum = (current_sum + curr[a]) % MOD;
        }
        prev_sum = current_sum; // Update total sum for the next iteration
        prev = curr;            // 'curr' becomes 'prev' for the next length
    }

    // After the loop, 'prev' (which was 'curr' in the last iteration)
    // contains the counts for arrays of length 'n'.
    return prev[x];
}

    Minor Optimization: Actually we don't even need 2 vector, since curr[a] only depends on prev[a]. You can actually just have 1 vector, and update in place.

    Stage 3: Full Optimization (O(N) Time, O(1) Space)

    Intuition: The core observation is that all numbers j from 2 to k behave identically in their recurrence relations, as long as j is not 1. This allows us to collapse k-1 distinct states into a single "ends with non-1" category. We only need to track two values:

    1. dp_ends_with_1: number of arrays of length L ending with 1.
    2. dp_ends_with_other: number of arrays of length L ending with any value in [2, k]

    Recurrence Relations:

    • To get next_dp_ends_with_1 (array ends with 1): The previous element must have been not 1. So, next_dp_ends_with_1 = dp_ends_with_other.
    • To get next_dp_ends_with_other (array ends with some j != 1):
      • If the previous element was 1: There were dp_ends_with_1 ways to end with 1. We can then append any of the (k-1) non-1 values. Contribution: dp_ends_with_1 * (k-1).
      • If the previous element was some p != 1: There were dp_ends_with_other ways to end with any non-1 value. We can then append any of the (k-2) non-1 values that are different from p. Contribution: dp_ends_with_other * (k-2).

    Code:

    long long countArray(int n, int k, int x) {
    long long MOD = 1000000007;

    // Handle base case where array length is 1
    if (n == 1) {
        return (x == 1) ? 1 : 0;
    }

    // dp_ends_with_1: Number of ways to form an array of length 'i' ending with 1
    long long dp_ends_with_1 = 0; 
    // dp_ends_with_other: Number of ways to form an array of length 'i' ending with a value != 1
    long long dp_ends_with_other = 0;

    // Initialize for arrays of length 1
    // The first element must be 1.
    dp_ends_with_1 = 1; // Array: [1]
    dp_ends_with_other = 0; // No way to end with other value if length is 1 and starts with 1

    // Iterate for lengths from 2 to n
    for (int i = 2; i <= n; ++i) {
        long long next_dp_ends_with_1;
        long long next_dp_ends_with_other;

        // To form an array of length 'i' ending with 1:
        // The previous element (at length i-1) must not be 1.
        // So, it comes from dp_ends_with_other.
        next_dp_ends_with_1 = dp_ends_with_other;
        
        // To form an array of length 'i' ending with a value j (where j != 1):
        // Case 1: Previous element was 1 (dp_ends_with_1 ways).
        // From these, we can choose any of the (k-1) values that are not 1.
        long long from_prev_1 = (dp_ends_with_1 * (k - 1)) % MOD;

        // Case 2: Previous element was some p (where p != 1) (dp_ends_with_other ways).
        // From these, we can choose any of the (k-2) non-1 values that are also not p.
        long long from_prev_other = (dp_ends_with_other * (k - 2)) % MOD;

        next_dp_ends_with_other = (from_prev_1 + from_prev_other) % MOD;

        // Update current states for the next iteration
        dp_ends_with_1 = next_dp_ends_with_1;
        dp_ends_with_other = next_dp_ends_with_other;
    }

    // After the loop, dp_ends_with_1 and dp_ends_with_other hold counts for length n.
    if (x == 1) {
        return dp_ends_with_1;
    } else {
        return dp_ends_with_other;
    }
}