Hi, in August, I also gave the OA for Uber Winter Internship and got the same exact question in OA, hence am sharing my approach. Hope it helps.
Approach: My approach was to build the solution step by step. The main issue with the brute force solution is that the number of edges can be O(N^2) which will give TLE. So let's try to solve an easier version of the problem.
First, let's forget about the 3 Dimensions, and let's try to solve the problem in 1 Dimension. Suppose you are given points like [x1, x2, x3 ... ] and so on. The cost of an edge is defined as |x1-x2|. Now think about how will approach this problem.
To minimize the cost, we can simply sort the points by X coordinates and our answer will be the sum of the difference of adjacent elements in the sorted array.
Visualize the same problem in terms of graphs. Considering each point as a node, we actually create a graph with edges as {x1, x2}, {x2, x3}, {x3, x4} and so on where x1 <= x2 <= x3 ... with the weight |x[i] - x[i+1]| and our answer is the MST of this graph.
We can extend the same problem to 2D and then to 3D. In 2D along with the edges in X direction, we will also consider edges in Y direction, so we sort locations by Y coordinate and create edges as {y1, y2}, {y2, y3} ... and so on with the weight |y[i] - y[i+1]|. Note that the coordinates are sorted by Y, so the minimum of {|x[i] - x[i+1]|, |y[i] - y[i+1]|} is |y[i] - y[i+1]|.
To extend it to 3D, we can do the same process with Z coordinates. So we will sort locations by Z coordinates and create edges like {z1, z2}, {z2, z3}, {z3, z4} and so on, with the weight as |z[i] - z[i+1]|.
Now we have a graph with N nodes and 3*N edges. Apply Kruskal MST to find the cost of MST for this graph which is our answer.
I was able to pass all the test cases with this approach.
#include<bits/stdc++.h>
using namespace std;
#define int long long
struct Edge {
int u, v, w;
};
// Use kruskal MST here.
int solve(vector<vector<int>> locations, int N) {
vector<Edge> edges;
vector<pair<int, int>> xCoordinates;
vector<pair<int, int>> yCoordinates;
vector<pair<int, int>> zCoordinates;
for(int i = 0; i < N; i ++) {
xCoordinates.push_back({locations[i][0], i});
yCoordinates.push_back({locations[i][1], i});
zCoordinates.push_back({locations[i][2], i});
}
sort(xCoordinates.begin(), xCoordinates.end());
sort(yCoordinates.begin(), yCoordinates.end());
sort(zCoordinates.begin(), zCoordinates.end());
vector<Edge> edgeList;
for(int i = 0; i + 1 < N; i ++) {
int weightX = locations[xCoordinates[i].second][3] - locations[xCoordinates[i + 1].second][3];
int weightY = locations[yCoordinates[i].second][3] - locations[yCoordinates[i + 1].second][3];
int weightZ = locations[zCoordinates[i].second][3] - locations[zCoordinates[i + 1].second][3];
edges.push_back({xCoordinates[i].second, weightX});
edges.push_back({yCoordinates[i].second, weightY});
edges.push_back({yCoordinates[i].second, weightZ});
}
// Find min cost using kruskal MST
int minCost = kruskalMST(edges);
}