136 lines
3.8 KiB
Zig
136 lines
3.8 KiB
Zig
const std = @import("std");
|
|
|
|
const Node = struct {
|
|
frequency: usize,
|
|
kind: union(enum) {
|
|
internal: struct {
|
|
left: *Node,
|
|
right: *Node,
|
|
},
|
|
leaf: u8,
|
|
},
|
|
|
|
fn initLeaf(frequency: usize, ch: u8) Node {
|
|
return .{
|
|
.frequency = frequency,
|
|
.kind = .{ .leaf = ch },
|
|
};
|
|
}
|
|
|
|
fn initInternal(
|
|
allocator: std.mem.Allocator,
|
|
left_child: Node,
|
|
right_child: Node,
|
|
) !Node {
|
|
const left = try allocator.create(Node);
|
|
const right = try allocator.create(Node);
|
|
left.* = left_child;
|
|
right.* = right_child;
|
|
return .{
|
|
.frequency = left_child.frequency + right_child.frequency,
|
|
.kind = .{ .internal = .{ .left = left, .right = right } },
|
|
};
|
|
}
|
|
|
|
fn deinit(self: Node, allocator: std.mem.Allocator) void {
|
|
switch (self.kind) {
|
|
.internal => |inner| {
|
|
inner.left.deinit(allocator);
|
|
inner.right.deinit(allocator);
|
|
allocator.destroy(inner.left);
|
|
allocator.destroy(inner.right);
|
|
},
|
|
.leaf => {},
|
|
}
|
|
}
|
|
|
|
fn compare(context: void, a: Node, b: Node) std.math.Order {
|
|
_ = context;
|
|
return std.math.order(a.frequency, b.frequency);
|
|
}
|
|
};
|
|
|
|
pub fn main() !void {
|
|
const text = "this is an example for huffman encoding";
|
|
|
|
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
|
|
defer std.debug.assert(gpa.deinit() == .ok);
|
|
|
|
const allocator = gpa.allocator();
|
|
var frequencies = std.AutoHashMap(u8, usize).init(allocator);
|
|
defer frequencies.deinit();
|
|
|
|
for (text) |ch| {
|
|
const gop = try frequencies.getOrPut(ch);
|
|
if (gop.found_existing) {
|
|
gop.value_ptr.* += 1;
|
|
} else {
|
|
gop.value_ptr.* = 1;
|
|
}
|
|
}
|
|
|
|
var prioritized_frequencies =
|
|
std.PriorityQueue(Node, void, Node.compare).init(allocator, {});
|
|
defer prioritized_frequencies.deinit();
|
|
|
|
var freq_it = frequencies.iterator();
|
|
while (freq_it.next()) |counted_char| {
|
|
try prioritized_frequencies.add(Node.initLeaf(
|
|
counted_char.value_ptr.*,
|
|
counted_char.key_ptr.*,
|
|
));
|
|
}
|
|
|
|
while (prioritized_frequencies.len > 1) {
|
|
try prioritized_frequencies.add(try Node.initInternal(
|
|
allocator,
|
|
prioritized_frequencies.remove(),
|
|
prioritized_frequencies.remove(),
|
|
));
|
|
}
|
|
|
|
const root = prioritized_frequencies.items[0];
|
|
defer root.deinit(allocator);
|
|
|
|
var codes = std.AutoArrayHashMap(u8, []const u8).init(allocator);
|
|
defer codes.deinit();
|
|
|
|
var arena = std.heap.ArenaAllocator.init(allocator);
|
|
defer arena.deinit();
|
|
|
|
try generateCodes(arena.allocator(), &root, &.{}, &codes);
|
|
|
|
const stdout = std.io.getStdOut().writer();
|
|
var code_it = codes.iterator();
|
|
while (code_it.next()) |item| {
|
|
try stdout.print("{c}: {s}\n", .{
|
|
item.key_ptr.*,
|
|
item.value_ptr.*,
|
|
});
|
|
}
|
|
}
|
|
|
|
fn generateCodes(
|
|
arena: std.mem.Allocator,
|
|
node: *const Node,
|
|
prefix: []const u8,
|
|
out_codes: *std.AutoArrayHashMap(u8, []const u8),
|
|
) !void {
|
|
switch (node.kind) {
|
|
.internal => |inner| {
|
|
const left_prefix = try arena.alloc(u8, prefix.len + 1);
|
|
std.mem.copy(u8, left_prefix, prefix);
|
|
left_prefix[prefix.len] = '0';
|
|
try generateCodes(arena, inner.left, left_prefix, out_codes);
|
|
|
|
const right_prefix = try arena.alloc(u8, prefix.len + 1);
|
|
std.mem.copy(u8, right_prefix, prefix);
|
|
right_prefix[prefix.len] = '1';
|
|
try generateCodes(arena, inner.right, right_prefix, out_codes);
|
|
},
|
|
.leaf => |ch| {
|
|
try out_codes.put(ch, prefix);
|
|
},
|
|
}
|
|
}
|