Why does my code TLE (Fenced In Gold)

USACO 2016 February Contest, Gold
Problem 3. Fenced In

My Work

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define mp make_pair
#define fi first
#define se second

int n, m, a[2010], b[2010], p[2100*2100];
ll ans = 0;
vector<pair<int,pair<int,int>>> edges;
int findSet(int i) { return (p[i]==i ? i : p[i] = findSet(p[i])); }
void unionSet(int i, int j) { p[findSet(j)]=findSet(i); }

int32_t main()
    ifstream cin("fencedin.in"); ofstream cout("fencedin.out");
    ios_base::sync_with_stdio(false); cin.tie(0);
    int A, B; cin >> A >> B >> n >> m; a[n+1]=A, b[m+1]=B, a[0] = 0, b[0] = 0;
    for(int i = 1; i <= n; i++) cin >> a[i];
    for(int i = 1; i <= m; i++) cin >> b[i];
    sort(a, a+n+2), sort(b, b+m+2);
    for(int i = 0; i < 2100*2090; i++) p[i]=i; // initialise dsu
    // add edges
    for(int i = 1; i <= m; i++)
        for(int j = 1; j <= n+1; j++)
            edges.push_back(mp(a[j]-a[j-1], mp(i*(n+1)+j, i*(n+1)+j+n+1)));
    for(int i = 1; i <= m+1; i++)
        for(int j = 1; j <= n; j++)
            edges.push_back(mp(b[i]-b[i-1], mp(i*(n+1)+j, i*(n+1)+j+1)));
    sort(edges.begin(), edges.end()); int num = 1; // kruskal
    for(int i = 0; i < (int)edges.size() and num<(n+1)*(m+1); i++)
            unionSet(edges[i].se.fi, edges[i].se.se), ans+=(ll)edges[i].fi, num++;
    cout << ans;

The code basically creates all edges and applies the standard kruskal’s algorithm to get minimum cost of building necessary edges(fences) to connect the entire graph(grid)

What I’ve Tried

I’ve tried optimizing this code e.g using fast input, breaking loop after creating enough edges, using long long only when necessary, but it still gives TLE on the last two testcases for some reason


I dont understand the reason for the TLE. Can anyone help me out?

Do union by size/depth?

I’ve tried union by size but it still TLEs.
One thing i noticed is the code with union by rank is faster by 200ms in the 8th testcase

The edges actually come in groups of equal weights. For example, the horizontal edges between x=2 and x=5 all have weight 3. Instead of sorting all of the edges, you could sort the weights of the groups of edges, which may be enough to pass.

1 Like

Thanks! It runs in about 1 sec at most for all cases now

Final Code: Fenced In.cpp