👌

zigで2次元のスライスを扱う

2023/11/05に公開

はじめに

最近zigで競プログラミングを始めました。競プロでは頻繁に2次元配列を使いますが、zigでは、どのように2次元配列を使えばよいか分からず困ったので記事を残すことにしました。

メモリ確保の方針

2次元配列を作成するとき、メモリをどのように確保するかは様々な方法があります。

次のリンクは、メモリ確保の方法について述べた素晴らしい記事です。

https://tondol.hatenablog.jp/entry/20090713/1247426321

競プロではより少ないメモリスペースで計算した方が有利なので、連続したメモリ空間を一括で割り当てて、それを2次元配列かのように使います。

たとえば 3×3の二次元配列の場合

const alc = std.heap.page_allocator;
var matrix = allocator.alloc(u8, 3 * 3)

のようにして一括で9マスを確保して、次の図のように使います。

添え字分かりづらい問題

連続したメモリを確保してあたかも2次元の配列かのように扱う場合、工夫をしないと添え字が分かりにくいという欠点があります。

例えば、2行3列にアクセスする場合、高級言語であればmatrix[1][2]のようにアクセスできますが、9マス分を一括でメモリ確保している場合、次のように書く必要があります。

// 2行3列にアクセス。colは行の長さ
matrix[1 * col + 2];

競プロ中は1bitでもワーキングメモリを節約したいので、このような形でアクセスするのは少し不便です。

改善しましょう。

工夫1 二次元配列を扱うためのsliceを用意する

次のような関数を定義します。

fn make2dSlice(alc: std.mem.Allocator, comptime T: type, row_size: usize, col_size: usize) ![][]T {
    const cntnt_bytes = try std.math.mul(usize, @sizeOf(T), row_size * col_size);
    const slice_bytes = try std.math.mul(usize, @sizeOf([]T), row_size);
    const total_bytes = try std.math.add(usize, cntnt_bytes, slice_bytes);
    const buf = try alc.alignedAlloc(u8, @alignOf([]T), total_bytes);
    errdefer alc.free(buf);

    const slice = std.mem.bytesAsSlice([]T, buf[0..slice_bytes]);
    for (0..row_size) |i| {
        const start = slice_bytes + i * col_size;
        const end = start + col_size;
        slice[i] = buf[start..end];
    }
    return slice;
}

すると、次のように高級言語と同じように書けます。

var matrix = try make2dSlice(allocator, u8, 3, 3);
// 2行3列にアクセス
matrix[1][2]

欠点はメモリ効率の悪さです。

厳密には違うと思いますが、上の関数が確保するメモリは次のようになっています。

配列に加えてスライスの分もメモリを確保しているため、生なやり方で配列を扱うときよりもメモリを余計に消費しています。(しかも行の数に比例して確保するメモリが増えていく。)

空間的な意味でメモリに余裕があるときはこれでもよいですが、そうではないケースもあります。

その場合は、次のような解決策があります。

工夫2 配列操作のための構造体を用意する。

次の関数を定義します。

// 謝辞
// このコードはSeedyROM(https://github.com/SeedyROM)さんがzigのコミュニティで
// 迷える私にくれたアドバイスをもとにしています。
// ありがとうございます。

pub fn Slice2d(comptime T: type, comptime _width: usize, comptime _height: usize) type {
    return struct {
        const Self = @This();
        pub const width: usize = _width;
        pub const height: usize = _height;

        data: []T,

        pub fn init(allocator: std.mem.Allocator) !Self {
            return .{
                .data = try allocator.alloc(T, width * height),
            };
        }

        pub fn deinit(self: *Self, allocator: std.mem.Allocator) void {
            allocator.free(self.data);
        }

        pub inline fn get(self: Self, x: usize, y: usize) T {
            return self.data[y * width + x];
        }

        pub inline fn set(self: *Self, x: usize, y: usize, value: T) void {
            self.data[y * width + x] = value;
        }
    };
}

次のように利用可能です。

const matrix = try Slice2d(u8, 2, 2).init(std.testing.allocator);
// 2行3列にアクセス
matrix.get(2, 3);

高級言語のmatrix[2][3]の記法よりは少し可読性が落ちるかもしませんが、十分実用的です。

さらに良いことに「工夫1」のメモリが余計に必要になる問題も解決しています。

さいごに

zigたのしい

Discussion