Was surpised to find Zig’s std doesn’t have a partition method, but I can write my own.
This is a very fucky algorithm that mashes a bunch of concepts together, but
at least it’s fast. I’ve let myself be told that Prims is a little faster still, but I’m proud of what I came up with on my own.
My part 1 has similar ideas, but with a modified UnionFind to keep track of the
cluster sizes.
part2:
fn solution(r: *Reader, _: Io, gpa: Allocator) !u64 {
// O(n)
const boxes: []const Box = blk: {
var boxes: ArrayList(Box) = .empty;
errdefer boxes.deinit(gpa);
while (try r.takeDelimiter('\n')) |line| {
var it = mem.tokenizeScalar(u8, line, ',');
const x: f32 = @floatFromInt(try fmt.parseInt(u32, it.next().?, 10));
const y: f32 = @floatFromInt(try fmt.parseInt(u32, it.next().?, 10));
const z: f32 = @floatFromInt(try fmt.parseInt(u32, it.next().?, 10));
assert(it.next() == null);
try boxes.append(gpa, .{ .x = x, .y = y, .z = z });
}
break :blk try boxes.toOwnedSlice(gpa);
};
defer gpa.free(boxes);
std.log.info("Found {d} boxes", .{boxes.len});
// (n * (n - 1)) / 2 == O(n^2)
var connections: ArrayList(Connection) = try .initCapacity(gpa, (boxes.len * (boxes.len - 1)) / 2);
defer connections.deinit(gpa);
for (0..boxes.len - 1) |i| {
for (i + 1..boxes.len) |j| {
connections.appendAssumeCapacity(.init(
boxes,
.from(@intCast(i)),
.from(@intCast(j)),
));
}
}
std.log.info("Found {d} connections", .{connections.items.len});
var uf: UnionFind = try .init(gpa, @intCast(boxes.len));
defer uf.deinit(gpa);
const connLessThanFn = struct {
pub fn connLessThanFn(_: void, lhs: Connection, rhs: Connection) bool {
return lhs.dist < rhs.dist;
}
}.connLessThanFn;
var clusters = boxes.len;
var stack: ArrayList([]Connection) = .empty;
defer stack.deinit(gpa);
// Partioning: avg O(nlogn) worst O(n^2)
// Sorting partitions: O(nlogn)
// Iterating conns: O(#connections) == O(n^2)
// Uniting: O(alpha(n))
//
// Strictly speaking O(n^2 * alpha(n))
try stack.append(gpa, connections.items);
while (stack.pop()) |slice| {
// cutoff is kinda arbitrary
if (slice.len > boxes.len) {
const pivot = partition(
Connection,
slice,
slice.len / 2,
{},
connLessThanFn,
);
// Put the latter half first, so the former half is popped first
try stack.append(gpa, slice[pivot..]);
try stack.append(gpa, slice[0..pivot]);
continue;
}
mem.sortUnstable(Connection, slice, {}, connLessThanFn);
for (slice) |connection| {
if (uf.unite(connection.id_a, connection.id_b)) {
clusters -= 1;
if (clusters == 1) {
const box_a = boxes[connection.id_a.to()];
const box_b = boxes[connection.id_b.to()];
const x_a: u64 = @intFromFloat(box_a.x);
const x_b: u64 = @intFromFloat(box_b.x);
return x_a * x_b;
}
}
}
}
unreachable;
}
pub fn partition(
comptime T: type,
items: []T,
partition_index: usize,
context: anytype,
comptime lessThanFn: fn (@TypeOf(context), lhs: T, rhs: T) bool,
) usize {
const swapIdx = struct {
pub fn swapIdx(list: []T, a: usize, b: usize) void {
if (a == b) return;
mem.swap(T, &list[a], &list[b]);
}
}.swapIdx;
assert(items.len >= 1);
assert(partition_index < items.len);
if (items.len == 1) return 0;
// Get the partition point to a known and stable location
swapIdx(items, partition_index, 0);
const partition_point = items[0];
var low_ptr: usize = 1;
var high_ptr: usize = items.len - 1;
while (low_ptr <= high_ptr) {
while (low_ptr <= high_ptr and
lessThanFn(context, items[low_ptr], partition_point)) : (low_ptr += 1)
{}
while (low_ptr <= high_ptr and
lessThanFn(context, partition_point, items[high_ptr])) : (high_ptr -= 1)
{}
if (low_ptr <= high_ptr) {
swapIdx(items, low_ptr, high_ptr);
low_ptr += 1;
high_ptr -= 1;
}
}
// swap the partition point into place
swapIdx(items, 0, low_ptr - 1);
return low_ptr - 1;
}
const Id = enum(u32) {
_,
pub fn from(i: u32) Id {
return @enumFromInt(i);
}
pub fn to(id: Id) u32 {
return @intFromEnum(id);
}
};
const UnionFind = struct {
const Element = struct {
parent: Id,
};
clusters: []Element,
pub fn init(gpa: Allocator, size: u32) !UnionFind {
const clusters = try gpa.alloc(Element, size);
for (0..size) |i|
clusters[i] = .{ .parent = .from(@intCast(i)) };
return .{ .clusters = clusters };
}
pub fn deinit(uf: UnionFind, gpa: Allocator) void {
gpa.free(uf.clusters);
}
pub fn find(uf: UnionFind, id: Id) Id {
const i = id.to();
const direct_parent = uf.clusters[i].parent;
if (direct_parent == id) {
return id;
}
const actual_parent = uf.find(direct_parent);
uf.clusters[i].parent = actual_parent; // 'cache' lookup
return actual_parent;
}
pub fn unite(uf: UnionFind, id_a: Id, id_b: Id) bool {
const parent_a = uf.find(id_a);
const parent_b = uf.find(id_b);
if (parent_a == parent_b) return false;
uf.clusters[parent_a.to()] = .{
.parent = parent_a,
};
uf.clusters[parent_b.to()] = .{
.parent = parent_a,
};
return true;
}
};
const Box = struct {
x: f32,
y: f32,
z: f32,
pub fn dist(a: Box, b: Box) f32 {
@setFloatMode(.optimized);
const square = struct {
pub fn square(e: f32) f32 {
return e * e;
}
}.square;
const res = @sqrt(square(a.x - b.x) + square(a.y - b.y) + square(a.z - b.z));
assert(res == 0 or math.isNormal(res));
assert(res >= 0);
return res;
}
};
const Connection = struct {
id_a: Id,
id_b: Id,
dist: f32,
pub fn init(boxes: []const Box, id_a: Id, id_b: Id) Connection {
const dist = boxes[id_a.to()].dist(boxes[id_b.to()]);
return .{ .id_a = id_a, .id_b = id_b, .dist = dist };
}
pub fn compare(a: Connection, b: Connection) math.Order {
return math.order(a.dist, b.dist);
}
};
these challenges are getting harder, I might struggle to finish the same day.