Submission #927133

#TimeUsernameProblemLanguageResultExecution timeMemory
927133vjudge1Paths (BOI18_paths)C++17
100 / 100
221 ms53840 KiB
/*
no more temmy :(
*/

#include<bits/stdc++.h>
using namespace std;
#pragma GCC optimize("O3,unroll-loops")
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp> 
using namespace __gnu_pbds;
// #include<icecream.hpp>
// using namespace icecream;
#define ll long long
#define int ll
#define ld long double
#define y1 cheza
// mt19937 rng(1983413);
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
template<class T> using ordered_set = tree<T,null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;
const int N=3e5+100;
const int M=1e3+1;
const int B=317;
const int mod=1e9+7;
const int INF=1e18;
const int lg=64;
const int dx[]={1,-1,0,0};
const int dy[]={0,0,1,-1};
const double eps=1e-9;
int n,m,k;
int c[N];
vector<int>g[N];
bool mark[N];
int ans=0;
void dfs(int x){
    mark[c[x]]=1;
    ans++;
    for(auto i:g[x]){
        if(!mark[c[i]]){
            dfs(i);
        }
    }
    mark[c[x]]=0;
}
void solve1(){
    for(int i=1;i<=n;i++){
        dfs(i);
    }
    cout<<ans-n<<'\n';
}
void solve2(){
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            if(c[j]!=c[i]){
                ans++;
            }
        }
    }
    cout<<ans<<'\n';
}   
int d[N][6];
void solve3(){
    for(int i=1;i<=n;i++){
        for(int j=1;j<=k;j++){
            d[i][j]=g[i].size();
        }
        for(auto j:g[i]){
            d[i][c[j]]--;
        }
    }
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            ans+=d[j][c[i]];
            ans++;
        }
    }
    cout<<ans<<'\n';
}
int d2[N][6];
void solve4(){
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            d[i][c[j]]++;
        }
    }
    for(int i=1;i<=n;i++){
        for(int l=1;l<=k;l++){
            for(auto j:g[i]){
                if(c[j]==l)continue;
                d2[i][l]+=(g[j].size()-d[j][c[i]]-d[j][l]);
            }
        }
    }
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            ans+=(g[j].size()-d[j][c[i]]);
            ans++;
        }
    }
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            int x=d2[j][c[i]];
            // (g[j].size()-d[j][c[i]])*(g[k].size()-d[k][c[j]]-d[k][c[i]]);
            ans+=x;
        }
    }
    cout<<ans<<'\n';
}
int d3[N][6][6];
void solve5(){
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            d[i][c[j]]++;
        }
    }
    for(int i=1;i<=n;i++){
        for(int l=1;l<=k;l++){
            for(auto j:g[i]){
                if(c[j]==l)continue;
                d2[i][l]+=(g[j].size()-d[j][c[i]]-d[j][l]);
            }
        }
    }
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            ans+=(g[j].size()-d[j][c[i]]);
            ans++;
        }
    }
    for(int i=1;i<=n;i++){
        for(auto j:g[i]){
            int x=d2[j][c[i]];
            // (g[j].size()-d[j][c[i]])*(g[k].size()-d[k][c[j]]-d[k][c[i]]);
            ans+=x;
        }
    }
    for(int i=1;i<=n;i++){
        for(int l=1;l<=k;l++){
            if(l==c[i])continue;
            for(int w=1;w<=k;w++){
                if(w==l||w==c[i])continue;
                for(auto j:g[i]){
                    if(w==c[j]||l==c[j])continue;
                    d3[i][l][w]+=(g[j].size()-d[j][l]-d[j][w]-d[j][c[i]]);
                }
            }
        }
    }
    // i,j,k,l,n
    // i,j,c[k],
    for(int i=1;i<=n;i++){
        // int cur=0;
        for(int l=1;l<=k;l++){
            if(l==c[i])continue;
            for(auto j:g[i]){
                if(c[j]==l)continue;
                ans+=d3[j][c[i]][l]*d[i][l];
                // cur+=d3[j][c[i]][l]*d[i][l];
            }
        }
        // cout<<cur<<' ';
    }
    cout<<ans<<'\n';
}
void test(){
    cin>>n>>m>>k;
    for(int i=1;i<=n;i++){
        cin>>c[i];
    }
    for(int i=1,x,y;i<=m;i++){
        cin>>x>>y;
        if(c[x]!=c[y]){
            g[x].push_back(y);
            g[y].push_back(x);
        }
    }
    if(k==1){
        cout<<0<<'\n';
        return;
    }
    if(k==2){
        solve2();
        return;
    }
    if(k==3){
        solve3();
        return;
    }
    if(k==4){
        solve4();
        return;
    }
    if(k==5){
        solve5();
        return;
    }
    if(n<=100&&m<=100&&k<=5){
        solve1();
        return;
    }
}
/*
*/  
 
signed main(){

    // ic.prefix("debug->| ");
    // freopen("input.txt","r",stdin);
    // freopen("output.txt","w",stdout);

    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    // cout.tie(nullptr);
    long long t2=1;
    //  cin>>t2;
    for(int i=1;i<=t2;i++){
        test();
    }
    
    return 0;
 
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...