DSA sheet / Dynamic programming
Array Description
Do these first: Dice Combinations, Grid Paths I
The problem in brief
You know an array of length n (up to 100 000) whose entries are between 1 and m (m up to 100), except that some entries are unknown (marked 0). Count the ways to fill the unknowns so that every pair of adjacent entries differs by at most 1. Print the count modulo 10^9 + 7.
An adapted, shortened summary (changed from the original) shared under the same licence, not a substitute for the statement. The problem is from the CSES Problem Set (Antti Laaksonen), CC BY-NC-SA 4.0. Read the official statement and submit at cses.fi/problemset/task/1746.
Try it
Tab indents; press Esc, then Tab, to leave the box. Ctrl+Enter runs.
Output
Errors
Hints
Stuck? Reveal one hint at a time. Each nudges without giving away the next.
-
Enumerating all fillings is m^n. What do you need to remember about a partial filling in order to keep extending it?
-
Only the last value matters for the next constraint. Define ways[i][v] as the number of valid fillings of the first i entries that end with value v.
-
Which values of the previous entry are compatible with v at the next position? How many are there?
-
Handle the known entries: at a known position only one value is allowed. How does that change the transition? And what is the answer at the end?
The walkthrough
Spoilers ahead: open a section only after you have given the hints a fair try.
How to think
The key is finding the right state. A partial filling has a lot of history, but the constraint “neighbours differ by at most 1” only ever looks at the last entry. So the state (position, last value) is a sufficient summary, and there are only n * m of those states, about 10^7.
Transition: ways[i][v] = ways[i - 1][v - 1] + ways[i - 1][v] + ways[i - 1][v + 1], but only if position i allows the value v (it is unknown, or it equals v); otherwise it is 0. Start: at position 0, ways[0][v] = 1 for every allowed v. Answer: the sum of ways[n - 1][v] over v.
The habit for counting constrained sequences: find the smallest summary of the past that determines which futures are legal (here just the last value), make that the DP dimension, and sweep position by position with a rolling array.
Intuition
Picture a token walking along the array positions and standing on a value at each step; it may stay put or move one up or one down, and at known positions it must stand on the given value. Count the walks by keeping, for each value, how many walks currently stand there.
Approach
- cur[v] for v in 1..m: if a[0] = 0, then 1 for every v, else 1 only for v = a[0].
- For each next position i: for every v allowed at i (all v if a[i] = 0, else only a[i]): next[v] = cur[v - 1] + cur[v] + cur[v + 1] (treat out-of-range as 0) mod 10^9 + 7; everything else is 0.
- The answer is the sum of cur[v] after the last position, mod 10^9 + 7.
Pitfalls: pad the arrays with an extra 0 slot on each side so v - 1 and v + 1 never go out of range; reduce modulo after summing the three terms (64-bit sum); a known value may make the total 0.
Complexity
O(n * m) time, O(m) memory.
Solutions
Written from scratch and checked by compiling and running each one against a brute force on random inputs. Fast input/output, the way you would submit it.
Show solutions (C++17, Python 3, Java 17, Node.js)
#include <bits/stdc++.h>
using namespace std;
int main() {
int n, m;
scanf("%d %d", &n, &m);
vector<int> a(n);
for (auto &v : a) scanf("%d", &v);
const long long MOD = 1000000007LL;
vector<long long> cur(m + 2, 0), nxt(m + 2, 0);
for (int v = 1; v <= m; v++) cur[v] = (a[0] == 0 || a[0] == v) ? 1 : 0;
for (int i = 1; i < n; i++) {
for (int v = 1; v <= m; v++) {
nxt[v] = (a[i] == 0 || a[i] == v) ? (cur[v - 1] + cur[v] + cur[v + 1]) % MOD : 0;
}
swap(cur, nxt);
}
long long total = 0;
for (int v = 1; v <= m; v++) total += cur[v];
printf("%lld\n", total % MOD);
return 0;
}import sys
def main():
data = sys.stdin.read().split()
n, m = int(data[0]), int(data[1])
a = list(map(int, data[2:2 + n]))
MOD = 10**9 + 7
cur = [0] * (m + 2)
for v in range(1, m + 1):
cur[v] = 1 if a[0] == 0 or a[0] == v else 0
for i in range(1, n):
nxt = [0] * (m + 2)
ai = a[i]
if ai == 0:
for v in range(1, m + 1):
nxt[v] = (cur[v - 1] + cur[v] + cur[v + 1]) % MOD
else:
nxt[ai] = (cur[ai - 1] + cur[ai] + cur[ai + 1]) % MOD
cur = nxt
print(sum(cur) % MOD)
main()import java.io.*;
public class Main {
public static void main(String[] args) throws IOException {
StreamTokenizer st = new StreamTokenizer(new BufferedInputStream(System.in));
st.nextToken();
int n = (int) st.nval;
st.nextToken();
int m = (int) st.nval;
int[] a = new int[n];
for (int i = 0; i < n; i++) {
st.nextToken();
a[i] = (int) st.nval;
}
final long MOD = 1_000_000_007L;
long[] cur = new long[m + 2];
long[] nxt = new long[m + 2];
for (int v = 1; v <= m; v++) cur[v] = (a[0] == 0 || a[0] == v) ? 1 : 0;
for (int i = 1; i < n; i++) {
for (int v = 1; v <= m; v++) {
nxt[v] = (a[i] == 0 || a[i] == v) ? (cur[v - 1] + cur[v] + cur[v + 1]) % MOD : 0;
}
long[] tmp = cur;
cur = nxt;
nxt = tmp;
}
long total = 0;
for (int v = 1; v <= m; v++) total += cur[v];
System.out.println(total % MOD);
}
}const data = require('fs').readFileSync(0, 'utf8').split(/\s+/).filter(Boolean).map(Number);
const n = data[0];
const m = data[1];
const a = data.slice(2, 2 + n);
const MOD = 1000000007;
let cur = new Array(m + 2).fill(0);
let nxt = new Array(m + 2).fill(0);
for (let v = 1; v <= m; v++) cur[v] = a[0] === 0 || a[0] === v ? 1 : 0;
for (let i = 1; i < n; i++) {
for (let v = 1; v <= m; v++) {
nxt[v] = a[i] === 0 || a[i] === v ? (cur[v - 1] + cur[v] + cur[v + 1]) % MOD : 0;
}
[cur, nxt] = [nxt, cur];
}
let total = 0;
for (let v = 1; v <= m; v++) total += cur[v];
console.log(total % MOD);Problem: CSES Problem Set (Antti Laaksonen, University of Helsinki), CC BY-NC-SA 4.0. The explanations and code above are original (© Anupam Kumar).