CSES Solutions – Array Description

You know that an array arr[] has N integers between 1 and M, and the absolute difference between two adjacent values is at most 1. Given a description of the array where some values may be unknown, your task is to count the number of arrays that match the description.

Examples:

Input: N = 3, M = 5, arr[] = {2, 0, 2}
Output: 3
Explanation: There are 3 arrays that match the description:

• First array: arr[] = {2, 1, 2}
• Second array: arr[] = {2, 2, 2}
• Third array: arr[] = {2, 3, 2}

Input: N = 1, M = 5, arr[] = {2}
Output: 1
Explanation: There is only one array that matches the description and that is {2}.

Approach: To solve the problem, follow the below idea:

The problem can be solved using Dynamic Programming. Maintain a dp[][] array, such that dp[i][j] stores the number of ways to have arr[i] = j.

Initially, let’s focus on the first index, i = 0. So, there are 2 possible values of arr[0].

• If arr[0] = 0, then we can have arr[0] as any value from 1 to M. So, for all values j from 1 to M, there is only one way to have arr[i] = j. Therefore, for all values of j from 1 to M, dp[0][j] = 1.
• If arr[0] != 0, then it means that arr[0] already has a value and we cannot put any other value at the first index. Therefore, we have dp[0][arr[0]] = 1.

Now, for all indices i from 1 to N, again there are two possible values of arr[i].

• If arr[i] = 0, then the number of ways we can have arr[i] = val is the sum of number of ways we can have arr[i – 1] = val – 1, arr[i-1] = val and arr[i – 1] = val + 1. Therefore, for all values of j from 1 to M, dp[i][j] = dp[i – 1][j – 1] + dp[i – 1][j] + dp[i – 1][j + 1].
• If arr[i] != 0, then it means that arr[i] already has a value and we cannot put any other value at index i. So, the number of ways will be sum of number of ways we can have arr[i – 1] = arr[i] – 1, arr[i-1] = arr[i] and arr[i – 1] = arr[i] + 1. Therefore, we have dp[i][arr[i]] = dp[i – 1][arr[i] – 1] + dp[i – 1][arr[i]] + dp[i – 1][arr[i] + 1].

The final answer will be sum of number of ways we can put any value j at the last index, that is sum of dp[N – 1][j] for all j from 1 to M.

Step-by-step algorithm:

• Declare a dp[][] array, such that dp[i][j] stores the number of ways to have arr[i] = j.
• Initialize all the values of dp[][] with 0.
• For index i = 0,
  • If arr[i] == 0, then for all values j from 1 to M, set dp[0][j] = 1.
  • Otherwise, set dp[0][arr[i]] = 1.
• For index i > 0,
  • If arr[i] == 0, then for all values j from 1 to M, set dp[i][j] = dp[i – 1][j – 1] + dp[i – 1][j] + dp[i – 1][j + 1].
  • Otherwise, set dp[i][arr[i]] = dp[i – 1][arr[i] – 1] + dp[i – 1][arr[i]] + dp[i – 1][arr[i] + 1]
• Since, the last index can have any value from 1 to M, for all j from 1 to M, sum up dp[N-1][j] to get the final answer.

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
#define ll long long int
#define mod 1000000007
#define INF 1000000000000
using namespace std;

// function to find number of possible arrays that satisfy
// the description
ll getans(vector<ll>& arr, ll N, ll M,
        vector<vector<ll> >& dp)
{
    // Iterate over all indices from 0 to N - 1
    for (int i = 0; i < N; i++) {
        if (i == 0) {
            // If we are at the first index and the value is
            // unknown
            if (arr[i] == 0) {
                // There is only one way to put any value at
                // the first index
                for (int val = 1; val <= M; val++)
                    dp[i][val] = 1;
            }
            else {
                // If we are at the first index and the
                // value is fixed, there is only one way to
                // put the value arr[0] at 0th index
                dp[i][arr[i]] = 1;
            }
        }
        else {
            // If we are not at the first index and the
            // value is unknown, then for each val from 1 to
            // M, the number of ways to put val are equal to
            // the number of ways we can put val-1, val,
            // val+1 at the previous index
            if (arr[i] == 0) {
                for (int val = 1; val <= M; val++) {
                    dp[i][val] = (dp[i - 1][val - 1]
                                + dp[i - 1][val]
                                + dp[i - 1][val + 1])
                                % mod;
                }
            }
            // If we are not at the first index and the
            // value is known, then the number of ways to
            // put arr[i] are equal to the number of ways we
            // can put arr[i]-1, arr[i], arr[i]+1 at the
            // previous index
            else {
                dp[i][arr[i]] = (dp[i - 1][arr[i] - 1]
                                + dp[i - 1][arr[i]]
                                + dp[i - 1][arr[i] + 1]) % mod;
            }
        }
    }

    // Variable to store the final answer
    ll ans = 0;

    // Sum the number of ways of putting any value in the
    // last index to get the final answer
    for (int val = 1; val <= M; val++) {
        ans = (ans + dp[N - 1][val]) % mod;
    }

    return ans;
}
int main()
{
    // Sample Input
    ll N = 3, M = 5;
    vector<ll> arr = {2, 0, 2};

    // dp[][] array such that dp[i][j] stores the number of
    // arrays that have arr[i] = j
    vector<vector<ll> > dp(N, vector<ll>(M + 2, 0));
    cout << getans(arr, N, M, dp) << endl;
}

Java

/*code by flutterfly */
import java.util.*;

public class Main {
    static long mod = 1000000007;
    static long INF = 1000000000000L;

    // Function to find the number of possible arrays that satisfy the description
    public static long getans(ArrayList<Long> arr, int N, int M, long[][] dp) {
        // Iterate over all indices from 0 to N - 1
        for (int i = 0; i < N; i++) {
            if (i == 0) {
                // If we are at the first index and the value is unknown
                if (arr.get(i) == 0) {
                    // There is only one way to put any value at the first index
                    for (int val = 1; val <= M; val++)
                        dp[i][val] = 1;
                } else {
                    // If we are at the first index and the value is fixed
                    // There is only one way to put the value arr[0] at 0th index
                    dp[i][arr.get(i).intValue()] = 1;
                }
            } else {
                // If we are not at the first index and the value is unknown
                if (arr.get(i) == 0) {
                    // Then for each val from 1 to M
                    // The number of ways to put val are equal to the number of ways
                    // we can put val-1, val, val+1 at the previous index
                    for (int val = 1; val <= M; val++) {
                        dp[i][val] = (dp[i - 1][val - 1]
                                + dp[i - 1][val]
                                + dp[i - 1][val + 1]) % mod;
                    }
                } else {
                    // If we are not at the first index and the value is known
                    // The number of ways to put arr[i] are equal to the number of ways
                    // we can put arr[i]-1, arr[i], arr[i]+1 at the previous index
                    dp[i][arr.get(i).intValue()] = (dp[i - 1][arr.get(i).intValue() - 1]
                            + dp[i - 1][arr.get(i).intValue()]
                            + dp[i - 1][arr.get(i).intValue() + 1]) % mod;
                }
            }
        }

        // Variable to store the final answer
        long ans = 0;

        // Sum the number of ways of putting any value in the
        // last index to get the final answer
        for (int val = 1; val <= M; val++) {
            ans = (ans + dp[N - 1][val]) % mod;
        }

        return ans;
    }

    public static void main(String[] args) {
        int N = 3, M = 5;
        // Sample Input
        ArrayList<Long> arr = new ArrayList<>(Arrays.asList(2L, 0L, 2L));

        // dp[][] array such that dp[i][j] stores the number of
        // arrays that have arr[i] = j
        long[][] dp = new long[N][M + 2];

        // Print the result
        System.out.println(getans(arr, N, M, dp));
    }
}

Python

# code by flutterfly
# Function to find the number of possible arrays that satisfy the description
def getans(arr, N, M, dp, mod):
    # Iterate over all indices from 0 to N - 1
    for i in range(N):
        if i == 0:
            # If we are at the first index and the value is unknown
            if arr[i] == 0:
                # There is only one way to put any value at the first index
                for val in range(1, M + 1):
                    dp[i][val] = 1
            else:
                # If we are at the first index and the value is fixed
                # There is only one way to put the value arr[0] at 0th index
                dp[i][arr[i]] = 1
        else:
            # If we are not at the first index and the value is unknown
            if arr[i] == 0:
                # Then for each val from 1 to M
                # The number of ways to put val are equal to the number of ways
                # we can put val-1, val, val+1 at the previous index
                for val in range(1, M + 1):
                    dp[i][val] = (dp[i - 1][val - 1] + dp[i - 1][val] + dp[i - 1][val + 1]) % mod
            else:
                # If we are not at the first index and the value is known
                # The number of ways to put arr[i] are equal to the number of ways
                # we can put arr[i]-1, arr[i], arr[i]+1 at the previous index
                dp[i][arr[i]] = (dp[i - 1][arr[i] - 1] + dp[i - 1][arr[i]] + dp[i - 1][arr[i] + 1]) % mod

    # Variable to store the final answer
    ans = 0

    # Sum the number of ways of putting any value in the
    # last index to get the final answer
    for val in range(1, M + 1):
        ans = (ans + dp[N - 1][val]) % mod

    return ans

if __name__ == "__main__":
    # Sample Input
    N = 3
    M = 5
    arr = [2, 0, 2]

    # Modulus value
    mod = 1000000007

    # dp[][] array such that dp[i][j] stores the number of
    # arrays that have arr[i] = j
    dp = [[0] * (M + 2) for _ in range(N)]

    # Print the result
    print(getans(arr, N, M, dp, mod))

C#

//codee by flutterfly
using System;
using System.Collections.Generic;

public class Program
{
    const long mod = 1000000007;
    const long INF = 1000000000000;

    // Function to find the number of possible arrays that satisfy the description
    public static long Getans(List<long> arr, long N, long M, long[][] dp)
    {
        // Iterate over all indices from 0 to N - 1
        for (int i = 0; i < N; i++)
        {
            if (i == 0)
            {
                // If we are at the first index and the value is unknown
                if (arr[i] == 0)
                {
                    // There is only one way to put any value at the first index
                    for (int val = 1; val <= M; val++)
                        dp[i][val] = 1;
                }
                else
                {
                    // If we are at the first index and the value is fixed
                    // There is only one way to put the value arr[0] at 0th index
                    dp[i][arr[i]] = 1;
                }
            }
            else
            {
                // If we are not at the first index and the value is unknown
                if (arr[i] == 0)
                {
                    // Then for each val from 1 to M
                    // The number of ways to put val are equal to the number of ways
                    // we can put val-1, val, val+1 at the previous index
                    for (int val = 1; val <= M; val++)
                    {
                        dp[i][val] = (dp[i - 1][val - 1]
                                    + dp[i - 1][val]
                                    + dp[i - 1][val + 1])
                                    % mod;
                    }
                }
                else
                {
                    // If we are not at the first index and the value is known
                    // The number of ways to put arr[i] are equal to the number of ways
                    // we can put arr[i]-1, arr[i], arr[i]+1 at the previous index
                    dp[i][arr[i]] = (dp[i - 1][arr[i] - 1]
                                    + dp[i - 1][arr[i]]
                                    + dp[i - 1][arr[i] + 1]) % mod;
                }
            }
        }

        // Variable to store the final answer
        long ans = 0;

        // Sum the number of ways of putting any value in the
        // last index to get the final answer
        for (int val = 1; val <= M; val++)
        {
            ans = (ans + dp[N - 1][val]) % mod;
        }

        return ans;
    }

    public static void Main(string[] args)
    {
        // Sample Input
        long N = 3, M = 5;
        List<long> arr = new List<long>() { 2, 0, 2 };

        // dp[][] array such that dp[i][j] stores the number of
        // arrays that have arr[i] = j
        long[][] dp = new long[N][];
        for (int i = 0; i < N; i++)
        {
            dp[i] = new long[M + 2];
        }

        // Print the result
        Console.WriteLine(Getans(arr, N, M, dp));
    }
}

Javascript

//code by flutterfly
// Function to find the number of possible arrays that satisfy the description
function getans(arr, N, M, dp) {
    const mod = 1000000007;
    // Iterate over all indices from 0 to N - 1
    for (let i = 0; i < N; i++) {
        if (i === 0) {
            // If we are at the first index and the value is unknown
            if (arr[i] === 0) {
                // There is only one way to put any value at the first index
                for (let val = 1; val <= M; val++)
                    dp[i][val] = 1;
            } else {
                // If we are at the first index and the value is fixed
                // There is only one way to put the value arr[0] at 0th index
                dp[i][arr[i]] = 1;
            }
        } else {
            // If we are not at the first index and the value is unknown
            if (arr[i] === 0) {
                // Then for each val from 1 to M
                // The number of ways to put val are equal to the number of ways
                // we can put val-1, val, val+1 at the previous index
                for (let val = 1; val <= M; val++) {
                    dp[i][val] = (dp[i - 1][val - 1]
                                + dp[i - 1][val]
                                + dp[i - 1][val + 1])
                                % mod;
                }
            } else {
                // If we are not at the first index and the value is known
                // The number of ways to put arr[i] are equal to the number of ways
                // we can put arr[i]-1, arr[i], arr[i]+1 at the previous index
                dp[i][arr[i]] = (dp[i - 1][arr[i] - 1]
                                + dp[i - 1][arr[i]]
                                + dp[i - 1][arr[i] + 1]) % mod;
            }
        }
    }

    // Variable to store the final answer
    let ans = 0;

    // Sum the number of ways of putting any value in the
    // last index to get the final answer
    for (let val = 1; val <= M; val++) {
        ans = (ans + dp[N - 1][val]) % mod;
    }

    return ans;
}

// Sample Input
const N = 3, M = 5;
const arr = [2, 0, 2];

// dp[][] array such that dp[i][j] stores the number of
// arrays that have arr[i] = j
const dp = Array.from(Array(N), () => new Array(M + 2).fill(0));

// Print the result
console.log(getans(arr, N, M, dp));

Output

3

Time Complexity: O(N * M), where N is the size of array arr[] and M is the range of values in arr[].
Auxiliary Space: O(N * M)