Skip to content

Instantly share code, notes, and snippets.

@Ryoga-exe

Ryoga-exe/DisjointSet.zig

Last active February 19, 2024 19:23
Show Gist options
  • Save Ryoga-exe/036533092aa23ee8f43b74121c63d7c0 to your computer and use it in GitHub Desktop.
Save Ryoga-exe/036533092aa23ee8f43b74121c63d7c0 to your computer and use it in GitHub Desktop.
Download ZIP
Zig implementation of disjoint set (union and find algorithm) and its tests.
Raw
DisjointSet.zig
const std = @import("std");
const Allocator = std.mem.Allocator;
const DisjointSet = struct {
const Self = @This();
roots: []usize,
sizes: []usize,
groups: usize,
allocator: Allocator,
pub fn init(allocator: Allocator, n: usize) Allocator.Error!Self {
var ret = Self{
.roots = try allocator.alloc(usize, n),
.sizes = try allocator.alloc(usize, n),
.groups = n,
.allocator = allocator,
};
for (0..n) |i| {
ret.roots[i] = i;
ret.sizes[i] = 1;
}
return ret;
}
pub fn deinit(self: *Self) void {
self.allocator.free(self.roots);
self.allocator.free(self.sizes);
}
pub fn find(self: *Self, index: usize) usize {
if (self.roots[index] == index) {
return index;
} else {
self.roots[index] = self.find(self.roots[index]);
return self.roots[index];
}
}
pub fn merge(self: *Self, u: usize, v: usize) bool {
var ru = self.find(u);
var rv = self.find(v);
if (ru == rv) {
return false;
}
if (self.sizes[ru] < self.sizes[rv]) {
const tmp = ru;
ru = rv;
rv = tmp;
}
self.roots[rv] = ru;
self.sizes[ru] += self.sizes[rv];
self.groups -= 1;
return true;
}
pub fn same(self: *Self, u: usize, v: usize) bool {
return self.find(u) == self.find(v);
}
pub fn len(self: Self) usize {
return self.roots.len;
}
};
test "AtCoder Library Practice Contest - A samples" {
const n = 4;
const query = [_]struct {
t: usize,
u: usize,
v: usize,
}{
.{ .t = 1, .u = 0, .v = 1 },
.{ .t = 0, .u = 0, .v = 1 },
.{ .t = 0, .u = 2, .v = 3 },
.{ .t = 1, .u = 0, .v = 1 },
.{ .t = 1, .u = 1, .v = 2 },
.{ .t = 0, .u = 0, .v = 2 },
.{ .t = 1, .u = 1, .v = 3 },
};
const expects = [_]?bool{
false,
null,
null,
true,
false,
null,
true,
};
var dsu = try DisjointSet.init(std.testing.allocator, n);
defer dsu.deinit();
for (query, expects) |q, e| {
if (q.t == 0) {
_ = dsu.merge(q.u, q.v);
} else {
try std.testing.expect(dsu.same(q.u, q.v) == e.?);
}
}
}
test "AtCoder Beginner Contest - D samples" {
const testcases = [_]struct {
n: usize,
m: usize,
p: []const usize,
x: []const usize,
y: []const usize,
}{
.{
.n = 5,
.m = 2,
.p = &[_]usize{ 5, 3, 1, 4, 2 },
.x = &[_]usize{ 1, 5 },
.y = &[_]usize{ 3, 4 },
},
.{
.n = 3,
.m = 2,
.p = &[_]usize{ 3, 2, 1 },
.x = &[_]usize{ 1, 2 },
.y = &[_]usize{ 2, 3 },
},
.{
.n = 10,
.m = 8,
.p = &[_]usize{ 5, 3, 6, 8, 7, 10, 9, 1, 2, 4 },
.x = &[_]usize{ 3, 4, 5, 2, 6, 3, 8, 7 },
.y = &[_]usize{ 1, 1, 9, 5, 5, 5, 9, 9 },
},
.{
.n = 5,
.m = 1,
.p = &[_]usize{ 1, 2, 3, 4, 5 },
.x = &[_]usize{1},
.y = &[_]usize{5},
},
};
const expects = [_]usize{
2,
3,
8,
5,
};
for (testcases, expects) |testcase, expect| {
var dsu = try DisjointSet.init(std.testing.allocator, testcase.n);
defer dsu.deinit();
for (testcase.x, testcase.y) |x, y| {
_ = dsu.merge(x - 1, y - 1);
}
var ans: usize = 0;
for (testcase.p, 0..) |p, i| {
ans += if (dsu.same(p - 1, i)) 1 else 0;
}
try std.testing.expect(ans == expect);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment