Expected Key Presses in the Klavir Problem

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.

Tags: KMP expectation Dynamic Programming Modular Arithmetic probabilistic analysis

Posted on Mon, 13 Jul 2026 16:09:08 +0000 by rslnerd