Golang - BitMap

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

package main

import (
	"errors"
)

// BitMap go 目前泛型还没有上线，我们定义 bitmap 只能固定一种类型，
// 这里用 uint64 就只是展示。添加了泛型后会更加灵活
type BitMap interface {
	// Add 添加一个元素到 bitmap 中
	Add(uint64) error
	// Del 从 bitmap 中删除一个元素
	Del(uint64) error
	// Contain bitmap 中是否包含这个数字
	Contain(uint64) bool
}

type bitmap struct {
	store []uint64
}

// Add 实现了 BitMap 接口的方法
func (b *bitmap) Add(u uint64) error {
	panic("implement me")
}

// Del 实现了 BitMap 接口的方法
func (b *bitmap) Del(u uint64) error {
	panic("implement me")
}

// Contain 实现了 BitMap 接口的方法
func (b *bitmap) Contain(u uint64) bool {
	panic("implement me")
}

// ErrZeroCapacity 表示要创建的 bitmap 的容量小于等于 0，这样是没有意义的
var ErrZeroCapacity = errors.New("capacity should greater than 0")

// NewBitmap 创建一个 Bitmap，cap 表示 bitmap 的容量
func NewBitmap(cap int) (*bitmap, error) {
	if cap <= 0 {
		return nil, ErrZeroCapacity
	}
	return &bitmap{
		store: make([]uint64, cap/64+1), // +1 防止数字太小整除后结果为 0
	}, nil
}

// 类型检查
var _ BitMap = (*bitmap)(nil)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

package main

import (
	"reflect"
	"testing"
)

func TestNewBitmap(t *testing.T) {
	type args struct {
		cap int
	}
	tests := []struct {
		name    string
		args    args
		want    *bitmap
		wantErr bool
	}{
		{
			name: "negative capacity",
			args: args{
				cap: -1,
			},
			want:    nil,
			wantErr: true,
		},
		{
			name: "ok",
			args: args{
				cap: 1,
			},
			want: &bitmap{
				store: make([]uint64, 1),
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := NewBitmap(tt.args.cap)
			if (err != nil) != tt.wantErr {
				t.Errorf("NewBitmap() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("NewBitmap() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func Test_BitMap(t *testing.T) {
	bm, _ := NewBitmap(10)
	if err := bm.Add(10); err != nil {
		t.Errorf("unexpected error: %v", err)
	}
	if !bm.Contain(10) {
		t.Errorf("bitmap should contain number 10")
	}

	if err := bm.Del(10); err != nil {
		t.Errorf("unexpect error: %s", err)
	}
	if bm.Contain(10) {
		t.Errorf("bitmap should not contain number 10")
	}
}

1
2
3
4

// 计算 u 在 bitmap 中的坐标
func (b *bitmap) coordinate(u uint64) (int, uint64) {
	return int(u / bitSize), 1 << (u % bitSize)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

// ErrOutOfRange 表示存储的元素超出了 bitmap 的最大容量
var ErrOutOfRange = errors.New("out of range")

// Add 实现了 BitMap 接口的方法
func (b *bitmap) Add(u uint64) error {
	x, y := b.coordinate(u)
	if x > len(b.store) {
		return ErrOutOfRange
	}
	b.store[x] |= y
	return nil
}

1
2
3
4
5
6
7
8
9

// Del 实现了 BitMap 接口的方法
func (b *bitmap) Del(u uint64) error {
	x, y := b.coordinate(u)
	if x > len(b.store) {
		return ErrOutOfRange
	}
	b.store[x] &= ^y
	return nil
}

1
2
3
4
5
6
7
8

// Contain 实现了 BitMap 接口的方法
func (b *bitmap) Contain(u uint64) bool {
	x, y := b.coordinate(u)
	if x > len(b.store) {
		return false
	}
	return b.store[x]&y == y
}

1
2

~/src/bitmap » go test .
ok      github.com/ningzio/bitmap       0.563s