Number of ways to Color the Boxes

Given N boxes (N >= 5) aligned in a line, we need to color all the boxes with using K colors such that the length of the longest segment of boxes with same color is equal to (N/2 + 1). The boxes in the segment should be contiguous. Find the total number of ways to color the boxes. Since the output may be very large, return the answer modulo 10⁹ + 7.

Examples:

Input: N = 5, K = 2
Output: 8
Explanation: Lets say we have two colors, C1 and C2. Then, we can color all the boxes in the following ways: [C1, C1, C1, C2, C1], [C1, C1, C1, C2, C2], [C2, C1, C1, C1, C2], [C2, C2, C1, C1, C1], [C2, C2, C2, C2, C1], [C2, C2, C2, C1, C1], [C1, C2, C2, C2, C1] and [C1, C1, C2, C2, C2].

Input: N = 5, K = 3
Output: 36
Explanation: There are 36 ways to color the boxes such that they follow the required conditions.

Approach:

The problem can be solved using Combinatorics and Binary Exponentiation. Since, we know the length of the segment = N/2+1, therefore there can be only one such segment. We can now divide our solution into 2 cases:

Case 1: When the segment is at one of the ends

• When the segment is at one of the ends, then we can color the box which is just adjacent to the segment with (K – 1) colors as it cannot have the same color as the segment. So for that box, we have (K – 1) choices to color and for the other remaining (N – N/2 – 2) boxes we can color each of them with any of the K colors. Also, we have 2 choices for the position of the segment, that is start and end. So, the total number of ways would be: (K-1) * K ^{(N – N/2 – 2)} * 2.

Case 2: When the segment is not at one of the ends

• When the segment is not at the ends, then we have one adjacent box on each side of the segment which we can color with (K-1) colors as they cannot have the same color as the segment. So for those two boxes, we have (K – 1) * (K – 1) choices to color and for the remaining (N – N/2 – 3) boxes, we can color each of them with any of the K colors. Also, we have (N – N/2 – 1 – 2) choices for the position of the segment. So, the total number of ways would be: (K – 1) * (K – 1) * K^{(N – N/2 – 3)} * (N – N/2 – 3)

Finally, we can sum up the two cases and multiply the sum by K as there are K choices to color the segment as well.

Step-by-step approach:

• Declare a function power(base, expo) to calculate the value of base^expo mod 10⁹ + 7
• Count the number of ways to color when the segment is at one of the ends: (K-1) * K ^{(N – N/2 – 2)} * 2
• Count the numbers of ways to color when the segment in not at one of the ends: (K – 1) * (K – 1) * K^{(N – N/2 – 3)} * (N – N/2 – 3)
• Return the product of K and sum of the above two cases as answer.

Below is the implementation of the above approach:

C++

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

const int MOD = 1e9 + 7;

// Method for binary exponentiation
ll power(ll base, ll expo)
{
    ll ans = 1;
    while (expo) {
        if (expo & 1) {
            ans = (ans * base) % MOD;
        }
        base = (base * base) % MOD;
        expo >>= 1LL;
    }
    return ans;
}

// Method to calculate the number of ways to color
// N boxes using K colors
ll countWays(ll N, ll K)
{
    // Case 1: When the segment is at one of the ends
    ll c1 = (((K - 1) * power(K, N - N / 2 - 2)) % MOD * 2)
            % MOD;
    // Case 2: When the segment is not at one of the ends
    ll c2 = ((((K - 1) * (K - 1)) % MOD
              * power(K, N - N / 2 - 3))
             % MOD * (N - N / 2 - 3))
            % MOD;
    // Multiplying the answer with K, as we have
    // K colors to color the segment
    return ((c1 + c2) % MOD * K) % MOD;
}

int main()
{
    // sample input
    ll N = 5, K = 2;

    // Function call
    cout << countWays(N, K) << "\n";

    return 0;
}

Java

import java.util.*;

public class Main {
    static final int MOD = 1000000007;

    // Method for binary exponentiation
    static long power(long base, long expo) {
        long ans = 1;
        while (expo > 0) {
            if ((expo & 1) == 1) {
                ans = (ans * base) % MOD;
            }
            base = (base * base) % MOD;
            expo >>= 1;
        }
        return ans;
    }

    // Method to calculate the number of ways to color
    // N boxes using K colors
    static long countWays(long N, long K) {
        // Case 1: When the segment is at one of the ends
        long c1 = (((K - 1) * power(K, N - N / 2 - 2) % MOD) * 2) % MOD;
        // Case 2: When the segment is not at one of the ends
        long c2 = ((((K - 1) * (K - 1) % MOD) *
                    power(K, N - N / 2 - 3) % MOD) * (N - N / 2 - 3)) % MOD;
        // Multiplying the answer with K, as we have
        // K colors to color the segment
        return ((c1 + c2) % MOD * K) % MOD;
    }

    public static void main(String[] args) {
        // Sample input
        long N = 5, K = 2;

        // Function call
        System.out.println(countWays(N, K));
    }
}
//Contributed by Aditi Tyagi

Python3

MOD = 1000000007

# Method for binary exponentiation
def power(base, expo):
    ans = 1
    while expo > 0:
        if expo & 1:
            ans = (ans * base) % MOD
        base = (base * base) % MOD
        expo >>= 1
    return ans

# Method to calculate the number of ways to color
# N boxes using K colors
def countWays(N, K):
    # Case 1: When the segment is at one of the ends
    c1 = (((K - 1) * power(K, N - N // 2 - 2) % MOD) * 2) % MOD
    # Case 2: When the segment is not at one of the ends
    c2 = ((((K - 1) * (K - 1) % MOD) *
           power(K, N - N // 2 - 3) % MOD) * (N - N // 2 - 3)) % MOD
    # Multiplying the answer with K, as we have
    # K colors to color the segment
    return ((c1 + c2) % MOD * K) % MOD

if __name__ == "__main__":
    # Sample input
    N = 5
    K = 2

    # Function call
    print(countWays(N, K))

C#

using System;

class GFG
{
    const int MOD = 1000000007;

    // Method for binary exponentiation
    static long Power(long baseNum, long expo)
    {
        long ans = 1;
        while (expo > 0)
        {
            if ((expo & 1) == 1)
            {
                ans = (ans * baseNum) % MOD;
            }
            baseNum = (baseNum * baseNum) % MOD;
            expo >>= 1;
        }
        return ans;
    }

    // Method to calculate the number of ways to color
    // N boxes using K colors
    static long CountWays(long N, long K)
    {
        // Case 1: When the segment is at one of the ends
        long c1 = (((K - 1) * Power(K, N - N / 2 - 2)) % MOD * 2) % MOD;

        // Case 2: When the segment is not at one of the ends
        long c2 = ((((K - 1) * (K - 1)) % MOD * Power(K, N - N / 2 - 3))
                   % MOD * (N - N / 2 - 3)) % MOD;

        // Multiplying the answer with K, as we have
        // K colors to color the segment
        return ((c1 + c2) % MOD * K) % MOD;
    }

    static void Main(string[] args)
    {
        // sample input
        long N = 5, K = 2;

        // Function call
        Console.WriteLine(CountWays(N, K));

        // Ensure the console stays open
        Console.ReadLine();
    }
}

Javascript

const MOD = 1e9 + 7;

// Method for binary exponentiation
function power(base, expo) {
    let ans = 1;
    while (expo) {
        if (expo & 1) {
            ans = (ans * base) % MOD;
        }
        base = (base * base) % MOD;
        expo >>= 1;
    }
    return ans;
}

// Method to calculate the number of ways to color
// N boxes using K colors
function countWays(N, K) {
    // Case 1: When the segment is at one of the ends
    let c1 = (((K - 1) * power(K, N - Math.floor(N / 2) - 2)) % MOD * 2) % MOD;
    // Case 2: When the segment is not at one of the ends
    let c2 = ((((K - 1) * (K - 1)) % MOD * power(K, N - Math.floor(N / 2) - 3)) % MOD * (N - Math.floor(N / 2) - 3)) % MOD;
    // Multiplying the answer with K, as we have K colors to color the segment
    return ((c1 + c2) % MOD * K) % MOD;
}

// Driver code

// sample input
let N = 5;
let K = 2;

// Function call
console.log(countWays(N, K));

Output

8

Time Complexity: log(K), where K is the number of colors available
Auxiliary Space: O(1)