Codeforces Round #286 (Div. 1) Problem D - Mr. Kitayuta's Colorful Graph

Problem Statement:
506D - Mr. Kitayuta's Colorful Graph

Solution:
The problem can be divided into (although not mutually exclusive) parts, first one is to effectively build the connectivity information in O(N), and the second part is to manage the queries. I am not so sure about the complexity analysis for the query handling part, but the following works.

First part: building the connectivity information using disjoint set union. Here is a simple yet powerful idea I learnt from the submissions made by clever gods who had solved the problem during the contest. We start out with an empty graph G. Then we incrementally add a new edge e. This edge has a color c and has two end points u and v. Before we add the edge and connect those vertices, we examine first whether each of u and v belong to any edges with color c. Let's say u indeed has such edge, call one of them e'. Since the addition of e to u will lead to e' and e being connected, we can simply represent this connection using a disjoint set union! So upon addition of e to u, e and e' are in the same set. Notice that we only need one of such edge e', since we have the invariant that every other e' incident to u will already be in the same disjoint set as e'! This is a very powerful observation, and this simplify the construction of the disjoint set union structure to O(N) overall.



Second part: handling the queries. For each query, we have a pair of vertices u and v. At this point we will iterate through the edges incident upon u, not all of them, but only one of each color (hence in essence we are iterating through all possible edge colors incident to u): For each such e with color c, we check if v has also an edge e' with color c (any one of them), which can be done in \(O(\lg{N})\) if we store this information using tree like structure. Then if e and e' belong to the same disjoint union set, we conclude that u is connected to v under color c. Hence the rest of the problem is just a matter of counting the number of such c. We can speed up the whole calculation by storing the computed results for every pair of u and v so far, and to pass the time limit, for each pair u,v being considered, we must only iterate through the vertex which is incident with least number of different colored edges.

I am thinking about the case where N(N+1)/2 = 100,000 = M, where every node has N-1 edges with different colors incident upon it. In this case, in each query (u,v), we must do \(O(N)\) checking, or \(O(N\lg{N})\) to be more precise, so for this case, N will be around 450, and the total operations needed will be something around 280 million computations, which is reasonable to be done in less than 4 seconds. In fact this will be the worst case running time, since if there are some nodes which has more than N-1 differently colored edge incident upon it, some nodes are necessarily have less of such edges, hence whenever possible for each (u,v) a query, we will want to choose the one with less number of edges. And I believe that the number of pairs with (u,v) both has more than N-1 number of differently colored edges will be considerably less. A better analysis with stronger observation might be best though..

#include <iostream>
#include <algorithm>
#include <vector>
#include <map>
#include <utility>
#include <cstdio>
using namespace std;

int par[100005];
int r[100005];
void init(){for(int i=0;i<=100000;++i)par[i]=i;}
int getpar(int u){return (par[u]==u?u:(par[u]=getpar(par[u])));}
void join(int u, int v){
    int x = getpar(u);
    int y = getpar(v);
    if(x == y) return;
    if(r[x] != r[y]){
        if(r[x] > r[y])par[y] = x;
        else par[x] = y;
    } else {
        r[x]++;
        par[y] = x;
    }
}
int N, M, Q;
int u,v,c;
map<int,int> color[100005];
map<pair<int,int>,int> ans;
typedef map<int,int>::iterator mit;
int main(){
    scanf("%d%d",&N,&M);
    init();
    for(int i=1;i<=M;++i){
        scanf("%d%d%d",&u,&v,&c);
        int a = 0, b = 0;
        if(color[u].count(c)) a = color[u][c];
        if(color[v].count(c)) b = color[v][c];
        color[u][c] = i;
        color[v][c] = i;
        if(a>0) join(i,a);
        if(b>0) join(i,b);
    }
    scanf("%d",&Q);
    while(Q--){
        scanf("%d%d",&u,&v);
        if(u>v)swap(u,v);
        if(ans.count(make_pair(u,v))){
            printf("%d\n",ans[make_pair(u,v)]);
            continue;
        }
        int ret = 0;
        if(color[u].size() > color[v].size())swap(u,v);
        for(mit uit = color[u].begin(); uit != color[u].end(); ++uit){
            mit vit = color[v].find(uit->first);
            if(vit == color[v].end()) continue;
            if(getpar(uit->second) == getpar(vit->second)) ++ret;
        }
        if(u>v)swap(u,v);
        ans[make_pair(u,v)] = ret;
        printf("%d\n", ret);
    }
    return 0;
}