前言

最近在看《编程珠玑》,开篇对于位图的使用让俺想起来之前面试时遇到的一个问题,假设直播间有大量用户,怎么快速的判断用户是否在线或离开?

现在想来,使用位图是不错的方式。首先用户ID是不重复的,创建BitSet(位图),用bit位0或1标识用户状态。

概念

BitSet(位图)是一组bit的集合。使用bit位置标识数字,bit状态0或1标识状态。

比如我们要存储数字集合 {1,2,3,5,8,13}

    13        8     5   3 2 1  
0 0 1 0 0 0 0 1 0 0 1 0 1 1 1 0

我们从右边开始数起(二进制左边为高位),将存在的数值设为1。

数据存储大小: 如1亿用户,只需要100000000 / 8 / 1024 / 1024 ≈ 12mb100000000/8/1024/1024≈12mb

设计目标

  • Add() 新增数据,可能会扩容
  • Clear() 清除数据,标识位由1变为0
  • Contains() 检查是否包含数据

数据结构

type BitSet struct {  
   data []uint64   
}

因为Golang没有bit数据类型。这里使用[]uint64存储数据。

BitSet

如果有bit类型,数据存储会更加简单。我们只需要定义bit的长度即可。类似:[]bit。

实现

预定义常量,BitSet 使用了比较多的位运算,位运算想比如常规数学运算更快些,当然也增加了阅读难度。

const (
	shift = 6    // 2^n = n of 64
	mask  = 0x3f // 2^6 - 1 = 63 = 0x3f
)

index函数定位数字在[]uint64的index。n » shift 等于 n / 64 posVal函数定位数字在uint64中的位置。 1 « uint(n&mask) 等于 1 * 2 ^ (n % 64)

func index(n int) int {
	return n >> shift
}

func posVal(n int) uint64 {
	return 1 << uint(n&mask)
}

NewBitSet

func NewBitSet(size int) *BitSet {
	return &BitSet{
		data: make([]uint64, size>>shift+1),
	}
}

size»shift+1等于size / shift + 1

Add

func (set *BitSet) Add(n int) *BitSet {
	if n < 0 {
		return set
	}

	i := index(n)
	if i >= len(set.data) {
		newData := make([]uint64, i+1)
		copy(newData, set.data)
		set.data = newData
	}

	if set.data[i]&posVal(n) == 0 {
		set.data[i] |= posVal(n)
	}
	return set
}

set.data[i]&posVal(n) == 0 为判断n是否存在于set.data[i]中,如果n存在,与的结果为posVal(n),否者为0。

set.data[i] |= posVal(n) 或预算,既将不存在的1位放到set.data[i]对应的二进制位中。

Clear

func (set *BitSet) Clear(n int) *BitSet {
	if n < 0 {
		return set
	}

	i := index(n)
	if i >= len(set.data) {
		return set
	}

	if set.data[i]&posVal(n) != 0 {
		set.data[i] &^= posVal(n)
	}
	return set
}

set.data[i] &^= posVal(n) 与亦或,相当于将已存在的posVal(n)清除。

Contains

func (set *BitSet) Contains(n int) bool {
	if n < 0 {
		return false
	}

	i := index(n)
	if i >= len(set.data) {
		return false
	}
	return set.data[i]&posVal(n) != 0
}

完整代码:https://github.com/mangoim/blog-code/tree/master/bitset

Reference