HackerEarth Array's force problem solution

In this HackerEarth Array'a force problem solution Let's consider some array A. The following algorithm calculates it's force:

Find all the continuous blocks where all the elemens of A are equal.

Calculate sum of squared lengths of these blocks.

For example if array A = {2, 3, 3, 1, 1, 1, 4, 2, 2} its force will be 12 + 22 + 32 + 12 + 22 = 19

We can reorder some elements and get array with greater force. In the example above we can move the first element of A to the end and get array A = {3, 3, 1, 1, 1, 4, 2, 2, 2} with force 22 + 32 + 12 + 32 = 23.

You are given an array. What is the maximum force you can get by reordering some of its elements?

HackerEarth Array's force problem solution.

#include <cstdio>
#include <algorithm>
#include <vector>
#include <iostream>
#include <string>
#include <string.h>
#include <cmath>
#include <set>
#include <map>
#include <bitset>
#include <iomanip>

#define X first
#define Y second
#define mp make_pair
#define pb push_back

typedef long long ll;

using namespace std;

const int MAXM = 1E6 + 100;
long long c[MAXM];
int n;
int a[MAXM], MOD;

void solve() {
    cin>>a[0]>>a[1]>>n>>MOD;
    for (int i = 0; i < MOD; i++) {
        c[i] = 0;
    }
    for (int i = 2; i < n; i++) {
        a[i] = a[i - 1] + a[i - 2];
        a[i] %= MOD;
    }
    for (int i = 0; i < n; i++) {
        //cerr<<a[i]<<" ";
        c[ a[i] ]++;
    }
    //cerr<<endl;
    long long ans = 0;
    for (int i = 0; i < MOD; i++) {
        ans += c[i] * c[i];
    }
    cout<<ans<<endl;
}

int main() {
    int t;
    cin>>t;
    while (t--) {
        solve();
    }
    return 0;
}

Second solution

#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<cctype>
#include<cstdlib>
#include<algorithm>
#include<bitset>
#include<vector>
#include<list>
#include<deque>
#include<queue>
#include<map>
#include<set>
#include<stack>
#include<cmath>
#include<sstream>
#include<fstream>
#include<iomanip>
#include<ctime>
#include<complex>
#include<functional>
#include<climits>
#include<cassert>
#include<iterator>
#include<unordered_set>
#include<unordered_map>
using namespace std;
int countt[1000001];
long long int p2[1000002];
namespace test {
    void end_test() {
        int val;
        if (cin >> val) {
            exit(1);
        }
    }
    void range_test(int t, int l, int r) {
        if (t < l || r < t) {
            exit(1);
        }
    }
}
int main()
{
    for(int i=0;i<1000002;i++){
        p2[i]=(long long int)(i)*(long long int)(i);
    }
    int t;
    scanf("%d",&t);
    test::range_test(t,1,100);
    while(t--){
        memset(countt,0,sizeof(countt));
        int a,b,c,d;
        scanf("%d%d%d%d",&a,&b,&c,&d);
        test::range_test(a,0,1000000);
        test::range_test(b,0,1000000);
        test::range_test(c,2,1000000);
        test::range_test(d,max(a,b)+1,1000000);
        countt[a]++;
        countt[b]++;
        register int k;
        c--;
        c--;
        while(c--){
            k=a+b;
            a=b;
            b=k;
            if(b>=d){
            b%=d;
            }
            countt[b]++;
        }
        long long int ans=0;
        for(int i=1000000;i>=0;i--){
            ans+=p2[countt[i]];
        }
        printf("%lld\n",ans);
    }
    test::end_test();
    return 0;
}