Notation
- Let \(dp_i\) denote the expected number of attempts to correctly play the prefix \(1 \to i\), with \(dp_1 = n\).
- The target melody is given by the array \(target_i\).
Analysis
When playing a melody, the current partially correct sequence falls into one of two cases:
- The current suffix matches some prefix of the melody.
- No suffix of the current attempt matches any prefix (the attempt is "independent").
The second case is simpler. For an independent attempt:
3
3
1 2 3
3
9
27
Now consider the first case, where a suffix matches a prefix. The matching prefixes can be precomputed with KMP.
Suppose we have correctly played up to \(i\) and want to play the \((i+1)\)-th note. Let the current pressed note be \(j\). If we press incorrectly, the attempt falls back to the longest already correct prefix; denote this length by \(k_j\). We have:
dp[1] = n;
for (int i = 1; i <= m; ++i) { // playing i-th note
long long sum = 0;
for (int j = 1; j <= n; ++j) { // possible pressed key
if (j != target[i + 1]) {
int k = fail[i];
while (k && target[k + 1] != j) k = fail[k];
if (target[k + 1] == j) ++k;
sum += dp[i] - dp[k];
}
}
dp[i + 1] = (sum + n + dp[i]) % mod;
}
Although acceptable, the nested loop can be optimised. Let \(f_{i,j}\) denote the fallback position when we are at position \(i\) and press note \(j\). Notice that \(f_{i,j}\) and \(f_{fail_i,j}\) differ only when \(j = target_{i+1}\). Hence the value of \(dp_{i+1}\) can inherit the answer from \(dp_{fail_{i+1}}\), adding back the contribution of the correct note \(j = target_{i+1}\). The contribution for the correct note is exactly \(n^{\,i+1}\).
Thus we obtain the following recurrence:
Code
#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9 + 7;
vector<int> compute_pi(const vector<int>& pattern) {
int m = pattern.size() - 1; // 1-indexed
vector<int> pi(m + 1);
for (int i = 2, j = 0; i <= m; ++i) {
while (j && pattern[j + 1] != pattern[i]) j = pi[j];
if (pattern[j + 1] == pattern[i]) ++j;
pi[i] = j;
}
return pi;
}
int main() {
int num_keys, len;
cin >> num_keys >> len;
vector<int> melody(len + 1);
for (int i = 1; i <= len; ++i) cin >> melody[i];
vector<int> pi = compute_pi(melody);
vector<int> result(len + 1);
long long power = num_keys % MOD;
for (int i = 1; i <= len; ++i) {
result[i] = (result[pi[i]] + power) % MOD;
cout << result[i] << "\n";
power = power * num_keys % MOD;
}
return 0;
}
Acknowledgements
Macesuted provided the optimisation idea.