888D - Almost Identity Permutations - CodeForces Solution

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

#define int long long

int ans = 0;

int nCr(int n, int r)
{
    int val = 1;

    for (int i = n; i >= n - r + 1; i--)
    {
        val *= i;
    }

    for (int i = 2; i <= r; i++)
    {
        val /= i;
    }

    return val;
}

void solve(int tc = 0)
{
    int n, k;
    cin >> n >> k;

    int cnt = 0;

    // The calculation of the derangement value of “N” is as follows
    // DN= N![1-1/1!+1/2!-1/3!…1/N!],

    // n = 2 Dn = 1
    // n = 3 Dn = 2
    // n = 4 Dn = 9

    if (k >= 2)
    {
        cnt += nCr(n, 2);
    }

    if (k >= 3)
    {
        // after selecting 3 elements
        // none of the 3 elements should be at their
        // correct position
        // let the elements be 4 5 6
        // so we can place 4 at
        // _ 4 _ which will become 6 4 5 or
        // _ _ 4 which will become 5 6 4
        cnt += 2 * nCr(n, 3);
    }

    if (k >= 4)
    {
        // to place the first element we have 3 options
        // now suppose we have placed the first element
        // at the incorrect position...now that leaves
        // us with one incorrect position (since first element
        // position is empty) and 2 correct positions
        // so we have 3 elements out of which we can place anyone
        // on the position of the first element and..and there is
        // only one way to arrange the other two...since both positions
        // are correct so we have total 3*3*1 = 9 ways
        cnt += 9 * nCr(n, 4);
    }

    cout << 1 + cnt << '\n';
}

int32_t main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int tc = 1;
    // cin >> tc;

    for (int i = 1; i <= tc; i++)
    {
        // cout << "Case #" << i << ": ";
        solve();
    }

    return 0;
}