• + 1 comment

    js

    const zero = BigInt(0);
    const one = BigInt(1);
    const two = BigInt(2);
    const mod = BigInt(10 ** 9 + 7);
    
    function modInverse(a) {
        let m = mod;
        let y = zero;
        let x = one;
    
        while (a > 1) {
            [a, m, y, x] = [m, a % m, x - a / m * y, y];
        }
    
        return x < 0 ? x + mod : x;
    }
    
    function modMul(...[a, ...args]) {
        return args.reduce((a, b) => (a * b) % mod, a);
    }
    
    function modAdd(...[a, ...args]) {
        return args.reduce((a, b) => (a + b) % mod, a);
    }
    
    function countArray(n, k, x) {
        k = BigInt(k);
        let last = 1 === x ? zero : one;
        let res = k - one - last;
        let total = modMul(k, k - one);
        const invK = modInverse(k);
    
        for (let i = 4; i <= n; i++) {
            [
                last,
                res,
                total
            ] = [
                    res,
                    modAdd(
                        total,
                        mod - modMul(two, total, invK),
                        last
                    ),
                    modMul(total, k - one)
                ];
        }
    
        return res;
    }