HackerRank Combinatorics – How Many Trees?

Challenge: How Many Trees?

Subdomeniu: Combinatorics (combinatorics)

Scor cont: 100.0 / 100

Submission status: Accepted

Submission score: 1.0

Submission ID: 464817523

Limbaj: cpp14

Link challenge: https://www.hackerrank.com/challenges/how-many-trees/problem

Cerinta

You are given two integers $N$ and $L$. Compute the number of unrooted labelled trees with exactly $N$ nodes and exactly $L$ leafs. This number can be huge, so please output it's modulo $10^9+7$.

Input Format

The first line of input consists of two space separated integers $N$ and $L$.

**Constraints**

$1 \leq L < N \leq 10^6$

Output Format

Display an answer to the problem on first line of the output.

Cod sursa

#include <bits/stdc++.h>

using namespace std;

/*
 * How Many Trees? (HackerRank)
 *
 * Count the number of labeled (1..n) undirected trees with exactly l leaves.
 * Answer modulo MOD (assumed 1e9+7 as in the original problem).
 *
 * Key fact (Prufer codes):
 *   Each labeled tree on n vertices <-> Prufer sequence of length n-2 over labels 1..n.
 *   Vertex v has degree deg(v) = 1 + (number of occurrences of v in the Prufer code).
 *   Leaves are exactly the labels that appear 0 times in the Prufer sequence.
 *
 * Therefore, a tree with exactly l leaves means exactly k = n-l labels appear at least once
 * in the code. Choose which l labels are missing, and then count surjections of length (n-2)
 * onto the remaining k labels.
 *
 * Count of surjections from an N-set onto k labels:
 *   k! * S(N,k) = sum_{i=0..k} (-1)^i * C(k,i) * (k-i)^N
 *
 * So final answer:
 *   ans = C(n,l) * sum_{i=0..k} (-1)^i C(k,i) (k-i)^(n-2)  (mod MOD)
 */

static const long long MOD = 1000000007LL;

static long long mod_pow(long long a, long long e) {
    long long r = 1 % MOD;
    a %= MOD;
    while (e > 0) {
        if (e & 1LL) r = (__int128)r * a % MOD;
        a = (__int128)a * a % MOD;
        e >>= 1LL;
    }
    return r;
}

static long long mod_inv(long long a) {
    return mod_pow(a, MOD - 2);
}

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

    long long n, l;
    if (!(cin >> n >> l)) return 0;

    if (n == 0) {
        cout << 0 << "\n";
        return 0;
    }

    if (n == 1) {
        // A single vertex has degree 0; by the standard definition "leaf = degree 1",
        // the only valid leaf count is 0.
        cout << ((l == 0) ? 1 : 0) << "\n";
        return 0;
    }

    if (l < 0 || l > n) {
        cout << 0 << "\n";
        return 0;
    }

    long long N = n - 2;       // length of Prufer code
    long long k = n - l;       // number of non-leaves (labels used at least once)

    if (k < 0) {
        cout << 0 << "\n";
        return 0;
    }

    // Precompute factorials up to n.
    int nn = (int)n;
    vector<long long> fact(nn + 1), invfact(nn + 1);
    fact[0] = 1;
    for (int i = 1; i <= nn; i++) fact[i] = (__int128)fact[i - 1] * i % MOD;
    invfact[nn] = mod_inv(fact[nn]);
    for (int i = nn; i >= 1; i--) invfact[i - 1] = (__int128)invfact[i] * i % MOD;

    auto C = [&](long long a, long long b) -> long long {
        if (b < 0 || b > a) return 0;
        return (long long)((__int128)fact[(int)a] * invfact[(int)b] % MOD * invfact[(int)(a - b)] % MOD);
    };

    long long choose_n_l = C(n, l);

    long long sum = 0;
    for (long long i = 0; i <= k; i++) {
        long long comb = C(k, i);
        long long base = (k - i) % MOD;
        long long pw;
        if (N == 0) {
            // treat 0^0 as 1
            pw = 1;
        } else {
            pw = mod_pow(base, N);
        }
        long long term = (__int128)comb * pw % MOD;
        if (i & 1LL) {
            sum -= term;
            if (sum < 0) sum += MOD;
        } else {
            sum += term;
            if (sum >= MOD) sum -= MOD;
        }
    }

    long long ans = (__int128)choose_n_l * sum % MOD;
    cout << ans << "\n";
    return 0;
}