Three colors triangles

I will assume that the formula in the link you provided is correct:

enter image description here


In order to avoid integer overflow, we will need to apply these modulo arithmetic rules:

(a * b) mod c = ((a mod c) * (b mod c)) mod c

(a ± b) mod c = ((a mod c) ± (b mod c)) mod c

Applying them to the formula:

enter image description here


But how do we calculate

enter image description here

... without directly calculating the coefficient itself (which can cause overflow)?

Since 3 is a prime number, this can be accomplished with Lucas's theorem:

enter image description here

... where n_i, m_i are the i-th digits of n, m in base-3.


Conversion to base-3 is easy with integer division:

// convert a number to base 3
// and returns the number of digits
unsigned conv_base_3(unsigned n, unsigned max, unsigned* out)
{
    unsigned i = 0;
    while (i < max && n > 0)
    {
        out[i] = n % 3;
        n /= 3;
        i++;
    }
    return i;
}

Note that since n_i, m_i are always in the range [0, 2] (because they are base-3 digits), C(n_i, m_i) are very easy to calculate:

// calculate the binomial coefficient for n < 3
unsigned binom_max_2(unsigned n, unsigned k)
{
    if (n < k)
        return 0;
    switch (n)
    {
        case 0:
        case 1:
            return 1;
        case 2:
            return 1 + (k == 1);

        // shouldn't happen
        default:
            return 0;
    }
}

And now the theorem itself:

// Lucas's theorem for p = 3
unsigned lucas_3(
    unsigned len_n, const unsigned * dig_n,
    unsigned len_k, const unsigned * dig_k
)
{
    // use modulo product rule:
    // prod[i] % 3 = ((prod[i - 1] % 3) * value[i])      
    unsigned prod = 1;
    for (unsigned i = 0; i < len_n; i++) {
        unsigned n_i = dig_n[i];
        unsigned k_i = (i < len_k) ? dig_k[i] : 0;
        prod = (prod * binom_max_2(n_i, k_i)) % 3;
    }
    return prod % 3;
}

Character conversion:

// convert from 012 to RGB
char int_2_char(int i)
{
    switch (i) {
        case 0: return 'R';
        case 1: return 'G';
        case 2: return 'B';

        // shouldn't happen
        default:
            return '\0';
    }
}

// convert from RGB to 012
unsigned char_2_int(char c)
{
    switch (c) {
        case 'R': return 0;
        case 'G': return 1;
        case 'B': return 2;

        // shouldn't happen
        default:
            return 3;
    }
}

Finally, the triangle algorithm:

// the problem constraints state that n <= 10 ** 5
// max number of base-3 digits
#define MAX_N_LOG_3 11

// main algorithm function
char triangle(const char * input)
{
    unsigned sum = 0;
    const int n = strlen(input);

    // calculate digits of n - 1
    unsigned dig_n[MAX_N_LOG_3];
    unsigned len_n = conv_base_3(n - 1, MAX_N_LOG_3, dig_n);

    for (unsigned km1 = 0; km1 < n; km1++)
    {
        // calculate digits of k - 1
        unsigned dig_k[MAX_N_LOG_3];
        unsigned len_k = conv_base_3(km1, MAX_N_LOG_3, dig_k);

        // calculate C(n - 1, k - 1) mod 3
        unsigned Cnk_mod3 = lucas_3(len_n, dig_n, len_k, dig_k);

        // add using the modulo rule
        sum = (sum + Cnk_mod3 * char_2_int(input[km1])) % 3;
    }

    // value of (-1) ** (n - 1)
    // (no need for pow; just need to know if n is odd or even)
    int sign = (n % 2) * 2 - 1;

    // for negative numbers, must resolve the difference
    // between C's % operator and mathematical mod
    int sum_mod3 = (3 + (sign * (int)(sum % 3)) % 3;
    return int_2_char(sum_mod3);
}

The above code passes all tests; note that it was written in favor of clarity, not performance.


So why was this code able to pass all tests in the allocated time, whereas the simple table-based approach wasn't? Because of its time complexity:

  • The table-based approach processes all levels of the triangle, which is O(n^2) (see Triangle Numbers).

  • Of course, using Lucas's algorithm, only the top level has to be processed; however the algorithm itself is O(log n), because it loops through every digit of n (regardless of the base). The overall complexity is O(n log n), which still represents a significant improvement.