实现一个 BitMap

楔子

本次来聊一聊 BitMap，这个在面试中经常会问到，在介绍它之前我们先来看一个问题。

假设当前你的服务有一千万个用户，每个用户都有一个整型ID，你要能够实时统计当前在线的人数，以及某个用户有没有登陆，这个时候你要怎么做呢？

首先我们想到用一个集合，来保存每个用户的 ID。当用户登陆时，就将 ID 插入到集合中；用户退出时，就将 ID 从集合中删除。所以集合中元素的数量就是当前的在线人数，而某个用户是否登陆就看其 ID 是否在集合中即可。

或者可以考虑使用数组，不是有一千万个用户吗？申请一个长度一千万的数组，初始时元素均为0，索引对应用户的ID。当用户登录时，就将该 ID 对应的数组元素设置为 1；用户退出时，将数组元素设置为 0。所以数组所有元素之和就是当前的在线人数，而某个用户是否登陆就看其 ID 对应的数组元素是否为 1 即可。

这是非常正确的思路，实现起来也非常方便，但是比较难缠的面试官可能会问你还有没有其它的做法。注意：并不是说这种做法不对，只是想顺着这个话题拓展一下思维。数组实现的话会耗费内存，如果活跃的用户非常多，那么该集合也会很耗费内存。这个时候面试官可能考察的就是 BitMap，下面就来介绍一下。

什么是 BitMap？

我们上面介绍了通过数组去存储，BitMap 也是类似的做法，只不过 BitMap 是用一个位表示一个用户。举个栗子，假设我们想表达 0 到 30，那么按照之前的逻辑，数组里面要有 31 个元素。但是采用 BitMap 的话，数组只需要一个元素即可，如果采用 long 的话，那么有 64 个位，足以表示这 31 个数。假设 ID 为 20 的用户登录了，我们就把该整数的第 21 个位设置为 1 即可。

那如果要表示 65 呢？一个 long 有 64 个位，最高能表示 63， 那么就再来一个整数，将第二个整数的第 3 个位设置为 1 即可。

所以原来数组中一个元素只能表示一个用户，现在一个元素可以表示 64 个用户，内存占用变成了原来的 \({\frac 1 {64}}\)。

实现 BitMap

下面我们就来使用 Python 实现一下 BitMap。

class BitMap:

    def __init__(self, max_count: int):
        """
        :param max_count: 最多容纳多少个用户
        """
        # 申请指定容量的列表，每个元素用 64 字节
        self.__array = [0] * ((max_count >> 6) + 1)

    def put_user(self, user_id: int):
        """
        将用户 ID 加入的 BitMap 中
        :param user_id:
        :return:
        """
        # 计算出所在的索引
        # 然后找到指定的位，将该位更新为 1
        self.__array[user_id >> 6] |= 1 << (user_id & 63)

    def get_user(self, user_id):
        """
        根据 user_id 获取用户，或者说判断用户是否登陆
        :param user_id:
        :return:
        """
        # 计算出所在的索引，判断指定的位是否为 1
        return self.__array[user_id >> 6] >> (user_id & 63) & 1 == 1

    def delete_user(self, user_id):
        """
        根据 user_id 删除用户
        :param user_id:
        :return:
        """
        self.__array[user_id >> 6] &= ~(1 << (user_id & 63))

    def online_user_count(self):
        """
        返回当前在线用户总数
        也就是统计每一个元素中 1 的位数，然后相加
        :return:
        """
        count = 0
        for item in self.__array:
            count += bin(item).count("1")
        return count


bm = BitMap(1000_0000)
print(bm.get_user(213456))  # False
bm.put_user(213456)
print(bm.get_user(213456))  # True
print(bm.online_user_count())  # 1
bm.delete_user(213456)
print(bm.get_user(213456))  # False
print(bm.online_user_count())  # 0

总的来说还算比较简单的。