HackerRank Probability – Colorful Polygon

Challenge: Colorful Polygon

Subdomeniu: Probability (probability)

Scor cont: 120.0 / 120

Submission status: Accepted

Submission score: 1.0

Submission ID: 464839985

Limbaj: cpp14

Link challenge: https://www.hackerrank.com/challenges/colorful-polygon/problem

Cerinta

You are given regular $n$-gon. Each vertex of the $n$-gon was randomly colored with one of $n$ colors. Your task is to find the number of special subsets of polygon vertices. Each special subset must meet all the requirements:

- The subset must contain at least two vertices.
- If vertices belonging to the subset are erased from the polygon (note, that the adjacent edges of those vertices will also get erased), the remaining vertices and edges form some continuous paths. None of those paths should contain two vertices of the same color.

Please, calculate the number of described special subsets and print it modulo $10^{9}+7$.

**Input Format**

The first line contains an integer $n$ $(3 \le n \le 10^5)$. The next line contains $n$ space-separated integers $a_1, a_2, \dots, a_n$ $(1 \le a_i \le n)$. Integer $a_i$ denotes the color of the $i$-th polygon vertex. All the vertices of the polygon are numbered from $1$ to $n$ in clockwise order.

You may assume that each $a_i$ was generated randomly, uniformly and independently from other values. For example, you may assume that $a_{i} = random(1, n)$ for each $i$. Where function $random(1, n)$ returns uniform integer from $1$ to $n$.

**Output Format**

Print a single integer $-$ the answer to the problem modulo $10^9+7$.

**Sample Input #1**

    4
    1 1 1 1

**Sample Output #1**

    7

**Sample Input #2**

    4
    4 2 3 1

**Sample Output #2**

    11

Cod sursa

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

static const long long MOD = 1000000007LL;

static inline long long modAdd(long long a, long long b) {
    a += b;
    if (a >= MOD) a -= MOD;
    return a;
}
static inline long long modSub(long long a, long long b) {
    a -= b;
    if (a < 0) a += MOD;
    return a;
}
static long long modPow(long long a, long long e) {
    long long r = 1 % MOD;
    a %= MOD;
    while (e > 0) {
        if (e & 1) r = (r * a) % MOD;
        a = (a * a) % MOD;
        e >>= 1;
    }
    return r;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    vector<int> c(N);
    for (int i = 0; i < N; i++) cin >> c[i];

    // If all colors are distinct, every deletion subset of size >=2 is valid.
    {
        vector<char> seen(N + 1, 0);
        bool allDistinct = true;
        for (int x : c) {
            if (seen[x]) { allDistinct = false; break; }
            seen[x] = 1;
        }
        if (allDistinct) {
            long long ans = modSub(modSub(modPow(2, N), N % MOD), 1);
            cout << ans << '\n';
            return 0;
        }
    }

    // Start index s (minimal deleted vertex) is valid only while prefix [0..s-1] has unique colors.
    int maxStart = N - 1;
    {
        vector<char> seen(N + 1, 0);
        for (int i = 0; i < N; i++) {
            int x = c[i];
            if (seen[x]) {
                maxStart = i; // valid s are 0..i
                break;
            }
            seen[x] = 1;
        }
    }

    vector<int> left(N, 1);
    vector<long long> dp(N + 1), pref(N + 1);

    vector<int> lastPos(N + 1, 0), mark(N + 1, 0);
    int stamp = 0;

    long long ans = 0;

    for (int s = 0; s <= maxStart; s++) {
        int M = N - 1;        // rotated length excluding s
        int P = N - s - 1;    // free prefix length (after s in original)

        ++stamp;
        int start = 1;
        for (int i = 1; i <= M; i++) {
            int col = c[(s + i) % N];
            if (mark[col] != stamp) {
                mark[col] = stamp;
                lastPos[col] = 0;
            }
            if (lastPos[col] >= start) start = lastPos[col] + 1;
            left[i] = start;
            lastPos[col] = i;
        }

        dp[0] = 1;
        pref[0] = 1;

        for (int i = 1; i <= P; i++) {
            int j = i - 1;
            int L = 0;
            if (j > 0) L = left[j] - 1;

            long long total = pref[i - 1];
            if (L > 0) total = modSub(total, pref[L - 1]);

            dp[i] = total;
            pref[i] = modAdd(pref[i - 1], dp[i]);
        }

        int Lfinal = left[M] - 1;
        if (Lfinal < 0) Lfinal = 0;

        long long totalWays = 0;
        if (Lfinal <= P) {
            totalWays = pref[P];
            if (Lfinal > 0) totalWays = modSub(totalWays, pref[Lfinal - 1]);
        }

        // remove subset deleting only s
        if (left[M] == 1) totalWays = modSub(totalWays, 1);

        ans = modAdd(ans, totalWays);
    }

    cout << ans << '\n';
    return 0;
}