New-year tree (Editorial)

Solution

Define a_i as the number of ways to assign colors such that the $i$-th color differs from the first color, and similarly, define b_i as the number of ways where the $i$-th color matches the first one. The answer is a_n.

To calculate both of these, we can use the following recursive formula:

while b_i = a_{i-1}. We can solve the recurrence for a_i to obtain the closed-form formula: a_i = (k-1)^i + (-1)^i \cdot (k-1). The cases when k = 1 or n = 1 are treated separately.

Time complexity: $O(n)$.

void solve(){
    int n, k;
    cin >> n >> k;
    if(n == 1){
        cout << k << '\n';
        return;
    }
    if(k == 1){
        cout << -1 << '\n';
        return;
    }
    ll a = 0, b = k;
    for(int i = 2; i <= n; i++){
        ll a1 = a * (k-2) + b * (k-1);
        ll b1 = a;
        a = a1;
        b = b1;
    }
    cout << (a == 0 ? -1 : a);
}