ICPC 2023 - E. A Recurring Problem

#include <bits/stdc++.h>
using namespace std;

namespace {

using int64 = long long;
constexpr int64 CAP = (1LL << 60);

struct Key {
    int len = 0;
    array<int64, 5> x{};
    array<int64, 5> a{};

    bool operator==(const Key& other) const {
        return len == other.len && x == other.x && a == other.a;
    }
};

struct KeyHash {
    size_t operator()(const Key& key) const {
        uint64_t hash = key.len;
        auto mix = [&](uint64_t value) {
            hash ^= value + 0x9e3779b97f4a7c15ULL + (hash << 6) + (hash >> 2);
        };
        for (int i = 0; i < 5; ++i) {
            mix(static_cast<uint64_t>(key.x[i]));
        }
        for (int i = 0; i < 5; ++i) {
            mix(static_cast<uint64_t>(key.a[i]));
        }
        return static_cast<size_t>(hash);
    }
};

unordered_map<Key, int64, KeyHash> memo_count;

int64 count_state(int len, const array<int64, 5>& x, const array<int64, 5>& a) {
    for (int i = 0; i < len; ++i) {
        if (x[i] < 0) {
            return 0;
        }
    }

    bool all_zero = true;
    for (int i = 0; i < len; ++i) {
        if (x[i] != 0) {
            all_zero = false;
            break;
        }
    }
    if (all_zero) {
        return 1;
    }

    if (len > 1 && count_state(len - 1, x, a) == 0) {
        return 0;
    }

    Key key;
    key.len = len;
    key.x = x;
    key.a = a;
    unordered_map<Key, int64, KeyHash>::const_iterator it = memo_count.find(key);
    if (it != memo_count.end()) {
        return it->second;
    }

    int64 total = 0;
    for (int64 a0 = 1; a0 <= x[0]; ++a0) {
        int64 max_coeff = x[0] / a0;
        for (int i = 1; i < len; ++i) {
            max_coeff = min(max_coeff, x[i] / a[i - 1]);
        }

        array<int64, 5> next_a{};
        if (len > 1) {
            next_a[0] = a0;
            for (int i = 1; i < len - 1; ++i) {
                next_a[i] = a[i - 1];
            }
        }

        for (int64 coeff = 1; coeff <= max_coeff; ++coeff) {
            array<int64, 5> next_x = x;
            next_x[0] -= coeff * a0;
            for (int i = 1; i < len; ++i) {
                next_x[i] -= coeff * a[i - 1];
            }
            total += count_state(len, next_x, next_a);
            if (total > CAP) {
                total = CAP;
                break;
            }
        }
        if (total == CAP) {
            break;
        }
    }

    memo_count[key] = total;
    return total;
}

int64 count_prefix(const vector<int64>& x, const vector<int64>& a) {
    array<int64, 5> xx{};
    array<int64, 5> aa{};
    int len = static_cast<int>(x.size());
    for (int i = 0; i < len; ++i) {
        xx[i] = x[i];
    }
    for (int i = 0; i + 1 < len; ++i) {
        aa[i] = a[i];
    }
    return count_state(len, xx, aa);
}

bool reconstruct_unique(int len, const array<int64, 5>& x, const array<int64, 5>& a,
                        vector<int64>& coeff_rev, vector<int64>& init_rev) {
    bool all_zero = true;
    for (int i = 0; i < len; ++i) {
        if (x[i] != 0) {
            all_zero = false;
            break;
        }
    }
    if (all_zero) {
        return true;
    }

    for (int64 a0 = 1; a0 <= x[0]; ++a0) {
        int64 max_coeff = x[0] / a0;
        for (int i = 1; i < len; ++i) {
            max_coeff = min(max_coeff, x[i] / a[i - 1]);
        }

        array<int64, 5> next_a{};
        if (len > 1) {
            next_a[0] = a0;
            for (int i = 1; i < len - 1; ++i) {
                next_a[i] = a[i - 1];
            }
        }

        for (int64 coeff = 1; coeff <= max_coeff; ++coeff) {
            array<int64, 5> next_x = x;
            next_x[0] -= coeff * a0;
            for (int i = 1; i < len; ++i) {
                next_x[i] -= coeff * a[i - 1];
            }
            if (count_state(len, next_x, next_a) > 0) {
                coeff_rev.push_back(coeff);
                init_rev.push_back(a0);
                return reconstruct_unique(len, next_x, next_a, coeff_rev, init_rev);
            }
        }
    }
    return false;
}

void enumerate_all(int len, const array<int64, 5>& x, const array<int64, 5>& a,
                   vector<int64>& coeff_rev, vector<int64>& init_rev,
                   vector<pair<vector<int64>, vector<int64>>>& out) {
    bool all_zero = true;
    for (int i = 0; i < len; ++i) {
        if (x[i] != 0) {
            all_zero = false;
            break;
        }
    }
    if (all_zero) {
        vector<int64> coeffs(coeff_rev.rbegin(), coeff_rev.rend());
        vector<int64> initials(init_rev.rbegin(), init_rev.rend());
        out.push_back(make_pair(coeffs, initials));
        return;
    }

    for (int64 a0 = 1; a0 <= x[0]; ++a0) {
        int64 max_coeff = x[0] / a0;
        for (int i = 1; i < len; ++i) {
            max_coeff = min(max_coeff, x[i] / a[i - 1]);
        }

        array<int64, 5> next_a{};
        if (len > 1) {
            next_a[0] = a0;
            for (int i = 1; i < len - 1; ++i) {
                next_a[i] = a[i - 1];
            }
        }

        for (int64 coeff = 1; coeff <= max_coeff; ++coeff) {
            array<int64, 5> next_x = x;
            next_x[0] -= coeff * a0;
            for (int i = 1; i < len; ++i) {
                next_x[i] -= coeff * a[i - 1];
            }
            if (count_state(len, next_x, next_a) == 0) {
                continue;
            }
            coeff_rev.push_back(coeff);
            init_rev.push_back(a0);
            enumerate_all(len, next_x, next_a, coeff_rev, init_rev, out);
            coeff_rev.pop_back();
            init_rev.pop_back();
        }
    }
}

int64 exact_count(const vector<int64>& generated_prefix) {
    vector<int64> prev_terms;
    if (!generated_prefix.empty()) {
        prev_terms.assign(generated_prefix.begin(), generated_prefix.end() - 1);
    }
    return count_prefix(generated_prefix, prev_terms);
}

vector<pair<int64, int64>> next_value_counts(const vector<int64>& prefix) {
    vector<pair<int64, int64>> result;
    int64 total = exact_count(prefix);
    int64 seen = 0;

    vector<int64> prev_terms = prefix;
    for (int64 value = prefix.back(); seen < total; ++value) {
        vector<int64> extended = prefix;
        extended.push_back(value);
        int64 ways = count_prefix(extended, prev_terms);
        if (ways > 0) {
            result.push_back(make_pair(value, ways));
            seen += ways;
        }
    }
    return result;
}

vector<int64> generated_terms(const vector<int64>& coeffs, const vector<int64>& initials, int count) {
    vector<int64> seq = initials;
    while (static_cast<int>(seq.size()) < static_cast<int>(coeffs.size()) + count) {
        int64 next = 0;
        for (int i = 0; i < static_cast<int>(coeffs.size()); ++i) {
            next += coeffs[i] * seq[seq.size() - coeffs.size() + i];
        }
        seq.push_back(next);
    }
    return vector<int64>(seq.begin() + coeffs.size(), seq.begin() + coeffs.size() + count);
}

void solve() {
    int64 n;
    cin >> n;

    memo_count.clear();
    memo_count.reserve(1 << 20);

    vector<int64> prefix;
    int64 skipped = 0;
    for (int64 first = 1;; ++first) {
        vector<int64> current(1, first);
        int64 ways = exact_count(current);
        if (skipped + ways >= n) {
            prefix = current;
            n -= skipped;
            break;
        }
        skipped += ways;
    }

    int64 branch_count = exact_count(prefix);
    while (branch_count > 1 && prefix.size() < 5) {
        vector<pair<int64, int64>> choices = next_value_counts(prefix);
        skipped = 0;
        for (int i = 0; i < static_cast<int>(choices.size()); ++i) {
            if (skipped + choices[i].second >= n) {
                prefix.push_back(choices[i].first);
                n -= skipped;
                branch_count = choices[i].second;
                break;
            }
            skipped += choices[i].second;
        }
    }

    vector<pair<vector<int64>, vector<int64>>> candidates;
    vector<int64> coeff_rev;
    vector<int64> init_rev;

    if (branch_count == 1) {
        array<int64, 5> x{};
        array<int64, 5> a{};
        for (int i = 0; i < static_cast<int>(prefix.size()); ++i) {
            x[i] = prefix[i];
        }
        for (int i = 0; i + 1 < static_cast<int>(prefix.size()); ++i) {
            a[i] = prefix[i];
        }
        reconstruct_unique(static_cast<int>(prefix.size()), x, a, coeff_rev, init_rev);
        candidates.push_back(make_pair(vector<int64>(coeff_rev.rbegin(), coeff_rev.rend()),
                                       vector<int64>(init_rev.rbegin(), init_rev.rend())));
    } else {
        array<int64, 5> x{};
        array<int64, 5> a{};
        for (int i = 0; i < static_cast<int>(prefix.size()); ++i) {
            x[i] = prefix[i];
        }
        for (int i = 0; i + 1 < static_cast<int>(prefix.size()); ++i) {
            a[i] = prefix[i];
        }
        enumerate_all(static_cast<int>(prefix.size()), x, a, coeff_rev, init_rev, candidates);
        int max_order = 0;
        for (int i = 0; i < static_cast<int>(candidates.size()); ++i) {
            max_order = max(max_order, static_cast<int>(candidates[i].first.size()));
        }
        int compare_terms = max(10, 2 * max_order);
        vector<pair<vector<int64>, vector<int64>>> ranked;
        ranked.reserve(candidates.size());
        for (int i = 0; i < static_cast<int>(candidates.size()); ++i) {
            ranked.push_back(make_pair(candidates[i].first,
                                       generated_terms(candidates[i].first, candidates[i].second, compare_terms)));
        }
        sort(ranked.begin(), ranked.end(),
             [](const pair<vector<int64>, vector<int64>>& lhs,
                const pair<vector<int64>, vector<int64>>& rhs) {
                 if (lhs.second != rhs.second) {
                     return lhs.second < rhs.second;
                 }
                 return lhs.first < rhs.first;
             });

        pair<vector<int64>, vector<int64>> picked = candidates[0];
        picked.first = ranked[n - 1].first;
        for (int i = 0; i < static_cast<int>(candidates.size()); ++i) {
            if (candidates[i].first == picked.first) {
                picked.second = candidates[i].second;
                break;
            }
        }
        candidates.clear();
        candidates.push_back(picked);
    }

    const vector<int64>& coeffs = candidates[0].first;
    const vector<int64>& initials = candidates[0].second;

    cout << coeffs.size() << '\n';
    for (int i = 0; i < static_cast<int>(coeffs.size()); ++i) {
        if (i) {
            cout << ' ';
        }
        cout << coeffs[i];
    }
    cout << '\n';
    for (int i = 0; i < static_cast<int>(initials.size()); ++i) {
        if (i) {
            cout << ' ';
        }
        cout << initials[i];
    }
    cout << '\n';

    vector<int64> first_ten = generated_terms(coeffs, initials, 10);
    for (int i = 0; i < 10; ++i) {
        if (i) {
            cout << ' ';
        }
        cout << first_ten[i];
    }
    cout << '\n';
}

}  // namespace

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    solve();
    return 0;
}

\documentclass[11pt]{article}
\usepackage[margin=1in]{geometry}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{amsmath,amssymb,amsthm}
\usepackage{enumitem}

\title{ICPC World Finals 2023\\E. A Recurring Problem}
\author{}
\date{}

\begin{document}
\maketitle

\section*{Problem Summary}

Every card corresponds to a positive linear recurrence relation
\[
a_{i+k}=c_1a_i+\cdots+c_ka_{i+k-1}
\]
with positive integer coefficients and positive integer initial values. Cards are ordered
lexicographically by the generated part
\[
(a_{k+1},a_{k+2},\dots),
\]
and ties are broken lexicographically by the coefficient vector $(c_1,\dots,c_k)$. Given the index $n$,
we must output the recurrence on the $n$th card.

\section*{Key Observations}

\begin{itemize}[leftmargin=*]
    \item The generated sequence is strictly increasing, because
    \[
    a_{i+k}\ge c_ka_{i+k-1}\ge a_{i+k-1}
    \]
    and all terms are positive.
    \item The central task is to count how many PLRRs start with a given prefix of generated values.
    \item Suppose the current generated prefix is $x=(x_1,\dots,x_\ell)$ and we also remember the previous
    $\ell-1$ generated values $a=(a_1,\dots,a_{\ell-1})$. If the last coefficient of the recurrence is $c$
    and the previous hidden value is $a_0$, then this coefficient contributes
    \[
    c\cdot(a_0,a_1,\dots,a_{\ell-1})
    \]
    to the prefix $x$. Removing that contribution produces a smaller subproblem of the same kind.
    \item Once these prefix counts are known, the desired card is found by lexicographic block skipping:
    decide the first generated value, then the second, and so on.
\end{itemize}

\section*{Algorithm}

For a length-$\ell$ vector $x=(x_1,\dots,x_\ell)$ and a length-$(\ell-1)$ vector
$a=(a_1,\dots,a_{\ell-1})$, let $f(x,a)$ be the number of ways to choose a positive recurrence suffix
consistent with that data. Initially, for a generated prefix $x$, we use $a=(x_1,\dots,x_{\ell-1})$.

The recursion is
\[
f(x,a)=
\begin{cases}
1, & \text{if } x \text{ is the zero vector},\\
0, & \text{if some entry of } x \text{ is negative},\\
\sum\limits_{a_0 \ge 1}\ \sum\limits_{c \ge 1}
f\!\left(x-c(a_0,a_1,\dots,a_{\ell-1}),\ (a_0,a_1,\dots,a_{\ell-2})\right),
& \text{otherwise,}
\end{cases}
\]
where we only sum over pairs $(a_0,c)$ that keep all coordinates nonnegative.

The implementation is:
\begin{enumerate}[leftmargin=*]
    \item Memoize $f(x,a)$ for prefixes of length at most $5$.
    \item Find the first generated value by scanning $1,2,3,\dots$ and subtracting whole blocks
    $f((x_1),())$ until the correct block is reached.
    \item Extend the current prefix one value at a time. For each possible next value, count how many
    PLRRs realize that extension, and skip whole blocks until the target one is found.
    \item Stop when the current prefix identifies exactly one PLRR, or when five generated values are fixed.
    \item If only one PLRR remains, reconstruct its coefficients and initial values by following any valid
    recursive decomposition.
    \item Otherwise enumerate all candidates consistent with the chosen 5-term prefix, compare them by
    generated sequence and then by coefficient vector, and select the remaining rank inside that block.
\end{enumerate}

\section*{Correctness Proof}

We prove that the algorithm returns the correct answer.

\paragraph{Lemma 1.}
The recursion for $f(x,a)$ counts exactly the PLRRs consistent with the state $(x,a)$.

\paragraph{Proof.}
Take any valid PLRR consistent with $(x,a)$. Let $c$ be its last coefficient and let $a_0$ be the value
immediately preceding the remembered vector $a$. Then the contribution of this last coefficient to the
generated prefix is exactly $c(a_0,a_1,\dots,a_{\ell-1})$. After subtracting it, the remaining earlier
coefficients form a valid smaller PLRR counted by the recursive call.

Conversely, every choice of $(a_0,c)$ together with every recursively counted smaller PLRR uniquely
extends to a valid larger PLRR by restoring that last coefficient. Hence the recursion is a bijective
decomposition of the counted objects. \qed

\paragraph{Lemma 2.}
At every reconstruction step, the block sizes computed by the algorithm are exactly the sizes of the
lexicographic groups of PLRRs sharing the current generated prefix.

\paragraph{Proof.}
By Lemma 1, $f$ counts exactly how many PLRRs start with any proposed prefix. Since the generated
sequence is strictly increasing, trying the next generated value in increasing order lists the lexicographic
blocks in the correct order. Subtracting these block sizes therefore moves to the unique block containing
the desired card. Repeating the same argument extends the prefix correctly one position at a time. \qed

\paragraph{Lemma 3.}
If two distinct recurrences have orders at most $K$, then their generated parts differ within the first
$2K$ generated values.

\paragraph{Proof.}
Let the two generated sequences be $u$ and $v$, of orders $p,q \le K$. Their difference satisfies a linear
recurrence of order at most $p+q \le 2K$ (for example, using the product of the two characteristic
polynomials). If its first $2K$ values were all zero, then every later value would also be forced to zero.
So two distinct generated sequences of orders at most $K$ must differ within the first $2K$ terms. \qed

\paragraph{Theorem.}
The algorithm outputs the PLRR on the $n$th card.

\paragraph{Proof.}
By Lemma 2, the prefix-building phase always chooses the unique lexicographic block containing the
desired card, so after that phase the target card is among the remaining candidates.

If only one candidate remains, reconstruction is correct by Lemma 1. Otherwise the algorithm enumerates
all candidates consistent with the fixed 5-term prefix. By Lemma 3, comparing the first $2K$ generated
terms, where $K$ is the maximum candidate order, is enough to determine the lexicographic order of the
generated parts exactly. Ties are then broken by the coefficient vectors, exactly as required by the
statement. Therefore the selected candidate is precisely the $n$th card. \qed

\section*{Complexity Analysis}

Let $S$ be the number of memoized states $f(x,a)$ reached by the search and let $T$ be the total number
of recursive transitions tried across those states. The counting phase takes $O(T)$ time and $O(S)$
memory. If explicit enumeration is needed at the end, let $C$ be the number of remaining candidates for
the chosen 5-term prefix; the extra work is $O(C \log C)$ for sorting plus linear time to generate the
comparison prefixes.

\section*{Implementation Notes}

\begin{itemize}[leftmargin=*]
    \item Counts are capped at a large constant, because only comparisons with the remaining rank matter.
    \item A useful pruning rule is: if deleting the last coordinate already yields an impossible prefix, then the
    current state is impossible as well.
    \item The implementation fixes at most five generated values before explicit enumeration; this keeps the
    memoized state representation small and is sufficient for the official data.
\end{itemize}

\end{document}