Count Vowels Permutation solution and explanation

def count_vowel_permutation(n):
    MOD = 10**9 + 7
    # Vowel states initialized
    dp = [1, 1, 1, 1, 1]  # Represents counts for a, e, i, o, u respectively
    for _ in range(1, n):
        dp_next = [0] * 5
        dp_next[0] = (dp[1]) % MOD  # a can follow e
        dp_next[1] = (dp[0] + dp[2]) % MOD  # e can follow a, i
        dp_next[2] = (dp[0] + dp[1] + dp[3] + dp[4]) % MOD  # i can follow a, e, o, u
        dp_next[3] = (dp[2] + dp[4]) % MOD  # o can follow i, u
        dp_next[4] = (dp[0]) % MOD  # u can follow a
        dp = dp_next
    return sum(dp) % MOD

# Usage: 
# print(count_vowel_permutation(1))  # Output: 5
# print(count_vowel_permutation(2))  # Output: 10

public class VowelPermutation {
    public int countVowelPermutation(int n) {
        int MOD = 1000000007;
        long[] dp = new long[]{1, 1, 1, 1, 1};  // a, e, i, o, u
        for (int i = 1; i < n; i++) {
            long[] dpNext = new long[5];
            dpNext[0] = (dp[1] + dp[2] + dp[4]) % MOD;  // a follows e, i, u
            dpNext[1] = (dp[0] + dp[2]) % MOD;          // e follows a, i
            dpNext[2] = (dp[1] + dp[3]) % MOD;          // i follows e, o
            dpNext[3] = dp[2] % MOD;                    // o follows i
            dpNext[4] = (dp[2] + dp[3]) % MOD;          // u follows i, o
            dp = dpNext;
        }
        long sum = 0;
        for (long x : dp) sum = (sum + x) % MOD;
        return (int) sum;
    }

    public static void main(String[] args) {
        VowelPermutation vp = new VowelPermutation();
        System.out.println(vp.countVowelPermutation(1));  // Output: 5
        System.out.println(vp.countVowelPermutation(2));  // Output: 10
    }
}

#include <iostream>
#include <vector>

class VowelPermutation {
public:
    int countVowelPermutation(int n) {
        const int MOD = 1000000007;
        std::vector<long> dp(5, 1); // a, e, i, o, u
        for (int i = 1; i < n; ++i) {
            std::vector<long> dpNext(5);
            dpNext[0] = (dp[1] + dp[2] + dp[4]) % MOD; // a follows e, i, u
            dpNext[1] = (dp[0] + dp[2]) % MOD;        // e follows a, i
            dpNext[2] = (dp[1] + dp[3]) % MOD;        // i follows e, o
            dpNext[3] = dp[2] % MOD;                  // o follows i
            dpNext[4] = (dp[2] + dp[3]) % MOD;        // u follows i, o
            dp = dpNext;
        }
        long sum = 0;
        for (long x : dp) sum = (sum + x) % MOD;
        return sum;
    }
};

int main() {
    VowelPermutation vp;
    std::cout << vp.countVowelPermutation(1) << std::endl; // Output: 5
    std::cout << vp.countVowelPermutation(2) << std::endl; // Output: 10
    return 0;
}

function countVowelPermutation(n) {
    const MOD = 10 ** 9 + 7;
    let dp = [1, 1, 1, 1, 1]; // a, e, i, o, u
    for (let i = 1; i < n; i++) {
        let dpNext = new Array(5).fill(0);
        dpNext[0] = (dp[1] + dp[2] + dp[4]) % MOD; // a follows e, i, u
        dpNext[1] = (dp[0] + dp[2]) % MOD;        // e follows a, i
        dpNext[2] = (dp[1] + dp[3]) % MOD;        // i follows e, o
        dpNext[3] = dp[2] % MOD;                  // o follows i
        dpNext[4] = (dp[2] + dp[3]) % MOD;        // u follows i, o
        dp = dpNext;
    }
    return dp.reduce((sum, x) => (sum + x) % MOD, 0);
}

// Example usage:
// console.log(countVowelPermutation(1)); // Output: 5
// console.log(countVowelPermutation(2)); // Output: 10

function countVowelPermutation(n: number): number {
    const f: number[] = Array(5).fill(1);
    const mod = 1e9 + 7;
    for (let i = 1; i < n; ++i) {
        const g: number[] = Array(5).fill(0);
        g[0] = (f[1] + f[2] + f[4]) % mod;
        g[1] = (f[0] + f[2]) % mod;
        g[2] = (f[1] + f[3]) % mod;
        g[3] = f[2];
        g[4] = (f[2] + f[3]) % mod;
        f.splice(0, 5, ...g);
    }
    return f.reduce((a, b) => (a + b) % mod);
}

func countVowelPermutation(n int) (ans int) {
	const mod int = 1e9 + 7
	f := make([]int, 5)
	for i := range f {
		f[i] = 1
	}
	for i := 1; i < n; i++ {
		g := make([]int, 5)
		g[0] = (f[1] + f[2] + f[4]) % mod
		g[1] = (f[0] + f[2]) % mod
		g[2] = (f[1] + f[3]) % mod
		g[3] = f[2] % mod
		g[4] = (f[2] + f[3]) % mod
		f = g
	}
	for _, x := range f {
		ans = (ans + x) % mod
	}
	return
}

const mod = 1e9 + 7;

function countVowelPermutation(n: number): number {
    const a: number[][] = [
        [0, 1, 0, 0, 0],
        [1, 0, 1, 0, 0],
        [1, 1, 0, 1, 1],
        [0, 0, 1, 0, 1],
        [1, 0, 0, 0, 0],
    ];
    const res = pow(a, n - 1);
    return res[0].reduce((a, b) => (a + b) % mod);
}

function mul(a: number[][], b: number[][]): number[][] {
    const [m, n] = [a.length, b[0].length];
    const c = Array.from({ length: m }, () => Array.from({ length: n }, () => 0));
    for (let i = 0; i < m; ++i) {
        for (let j = 0; j < n; ++j) {
            for (let k = 0; k < b.length; ++k) {
                c[i][j] =
                    (c[i][j] + Number((BigInt(a[i][k]) * BigInt(b[k][j])) % BigInt(mod))) % mod;
            }
        }
    }
    return c;
}

function pow(a: number[][], n: number): number[][] {
    let res: number[][] = [[1, 1, 1, 1, 1]];
    while (n) {
        if (n & 1) {
            res = mul(res, a);
        }
        a = mul(a, a);
        n >>>= 1;
    }
    return res;
}

const mod = 1e9 + 7

func countVowelPermutation(n int) (ans int) {
	a := [][]int{
		{0, 1, 0, 0, 0},
		{1, 0, 1, 0, 0},
		{1, 1, 0, 1, 1},
		{0, 0, 1, 0, 1},
		{1, 0, 0, 0, 0}}
	res := pow(a, n-1)
	for _, x := range res[0] {
		ans = (ans + x) % mod
	}
	return
}

func mul(a, b [][]int) [][]int {
	m, n := len(a), len(b[0])
	c := make([][]int, m)
	for i := range c {
		c[i] = make([]int, n)
	}
	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			for k := 0; k < len(b); k++ {
				c[i][j] = (c[i][j] + a[i][k]*b[k][j]) % mod
			}
		}
	}
	return c
}

func pow(a [][]int, n int) [][]int {
	res := [][]int{{1, 1, 1, 1, 1}}
	for n > 0 {
		if n&1 == 1 {
			res = mul(res, a)
		}
		a = mul(a, a)
		n >>= 1
	}
	return res
}