USACO 2014 January Contest, Silver Problem 2. Cross Country

#include <bits/stdc++.h>
using namespace std;
#define Mx 510
struct Edge
{
    int ux,uy,vx,vy,w,val;
    Edge() {}
    Edge(int _ux,int _uy,int _vx,int _vy,int _w)
    {
        ux = _ux;
        uy = _uy;
        vx = _vx;
        vy = _vy;
        w = _w;
    }
};
int grid[Mx][Mx];
bool f[Mx][Mx];
vector<Edge>Edges;
int parent[Mx][Mx];
int N,M;
void init()
{
    memset(grid,0,sizeof grid);
    memset(f,0,sizeof f);
    memset(parent,0,sizeof parent);
    Edges.clear();
}


bool comp(Edge e1,Edge e2)
{
    return e1.w < e2.w;
}

int find_parent(int x,int y)
{
    int c = x*1000 + y;
    if(parent[x][y]==c)return c;
    int p = parent[x][y];
    return parent[x][y] = find_parent(p/1000,p%1000);
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    while(cin >> N >> M)
    {
        init();
        for(int i=0; i<N; i++)
        {
            for(int j=0; j<M; j++)
            {
                cin >> grid[i][j];
                parent[i][j] = i*1000 + j;
            }
        }
        int total_one = 0;

        for(int i=0; i<N; i++)
        {
            for(int j=0; j<M; j++)
            {
                cin >> f[i][j];
                if(f[i][j]==1)total_one++;
            }
        }

        int dx[]= {0,0,1,-1};
        int dy[]= {1,-1,0,0};

        for(int i=0; i<N; i++)
        {
            for(int j=0; j<M; j++)
            {
                for(int k=0; k<4; k++)
                {
                    int nx = i + dx[k];
                    int ny = j + dy[k];
                    if(nx>=0 && nx<N && ny>=0 && ny<M)
                    {
                        Edge e = Edge(i,j,nx,ny,abs(grid[i][j]-grid[nx][ny]));
                        e.val = f[i][j] || f[nx][ny];
                        Edges.push_back(e);

                    }
                }
            }
        }

        set<pair<int,int> >SET;
        SET.clear();
        sort(Edges.begin(),Edges.end(),comp);
        int mx = 0;
        int i=0;
        set<int>parents;
        for(i=0; i<Edges.size(); i++)
        {
            Edge e = Edges[i];
            int px = find_parent(e.ux,e.uy);
            int py = find_parent(e.vx,e.vy);
            if(px!=py)
            {
                if(parents.find(px)!=parents.end())
                {
                    parent[py/1000][py%1000] = px;
                    auto idx = parents.find(py);
                    if(idx!=parents.end())parents.erase(idx);
                }
                else if(parents.find(py)!=parents.end())
                {
                    parent[px/1000][px%1000] = py;
                }
                else
                {
                    parent[px/1000][px%1000] = py;
                    if(f[e.ux][e.uy] || f[e.vx][e.vy])
                    {
                        parents.insert(py);
                    }
                }
            }
            mx = max(mx,e.w);
            if(f[e.ux][e.uy])
            {
                SET.insert(make_pair(e.ux,e.uy));
            }
            if(f[e.vx][e.vy])
            {
                SET.insert(make_pair(e.vx,e.vy));

            }
            if(SET.size()==total_one && parents.size()==1)break;
        }
        cout << mx << endl;
    }

}