【力扣】设计内存分配器(高效实现)

题目

给你一个整数 n ，表示下标从 0 开始的内存数组的大小。所有内存单元开始都是空闲的。

请你设计一个具备以下功能的内存分配器：

分配 一块大小为 size 的连续空闲内存单元并赋 id mID 。
释放 给定 id mID 对应的所有内存单元。
注意：

多个块可以被分配到同一个 mID 。
你必须释放 mID 对应的所有内存单元，即便这些内存单元被分配在不同的块中。
实现 Allocator 类：

Allocator(int n) 使用一个大小为 n 的内存数组初始化 Allocator 对象。
int allocate(int size, int mID) 找出大小为 size 个连续空闲内存单元且位于 最左侧 的块，分配并赋 id mID 。返回块的第一个下标。如果不存在这样的块，返回 -1 。
int free(int mID) 释放 id mID 对应的所有内存单元。返回释放的内存单元数目。

示例：

输入
[“Allocator”, “allocate”, “allocate”, “allocate”, “free”, “allocate”, “allocate”, “allocate”, “free”, “allocate”, “free”]
[[10], [1, 1], [1, 2], [1, 3], [2], [3, 4], [1, 1], [1, 1], [1], [10, 2], [7]]
输出
[null, 0, 1, 2, 1, 3, 1, 6, 3, -1, 0]

解释

Allocator loc = new Allocator(10); // 初始化一个大小为 10 的内存数组，所有内存单元都是空闲的。
loc.allocate(1, 1); // 最左侧的块的第一个下标是 0 。内存数组变为 [1, , , , , , , , , ]。返回 0 。
loc.allocate(1, 2); // 最左侧的块的第一个下标是 1 。内存数组变为 [1,2, , , , , , , , ]。返回 1 。
loc.allocate(1, 3); // 最左侧的块的第一个下标是 2 。内存数组变为 [1,2,3, , , , , , , ]。返回 2 。
loc.free(2); // 释放 mID 为 2 的所有内存单元。内存数组变为 [1, ,3, , , , , , , ] 。返回 1 ，因为只有 1 个 mID 为 2 的内存单元。
loc.allocate(3, 4); // 最左侧的块的第一个下标是 3 。内存数组变为 [1, ,3,4,4,4, , , , ]。返回 3 。
loc.allocate(1, 1); // 最左侧的块的第一个下标是 1 。内存数组变为 [1,1,3,4,4,4, , , , ]。返回 1 。
loc.allocate(1, 1); // 最左侧的块的第一个下标是 6 。内存数组变为 [1,1,3,4,4,4,1, , , ]。返回 6 。
loc.free(1); // 释放 mID 为 1 的所有内存单元。内存数组变为 [ , ,3,4,4,4, , , , ] 。返回 3 ，因为有 3 个 mID 为 1 的内存单元。
loc.allocate(10, 2); // 无法找出长度为 10 个连续空闲内存单元的空闲块，所有返回 -1 。
loc.free(7); // 释放 mID 为 7 的所有内存单元。内存数组保持原状，因为不存在 mID 为 7 的内存单元。返回 0 。

提示：

1 <= n, size, mID <= 1000
最多调用 allocate 和 free 方法 1000 次

解题思路

因为数据量不大，可以直接用数组来做，但这里提供另一种高效一些的实现方式：

我们可以用起始地址+大小的结构来表示一个内存块:

struct Mem{
    int addr;
    int size;
};

同时优先使用地址更左侧的、大小合适的内存块，所以还提供一下比较器：

inline bool operator<(const Mem &m1, const Mem &m2){
    if(m1.addr != m2.addr)
        return m1.addr < m2.addr;
    return m1.size < m2.size;
}

分配器在多次分配释放操作后可能会有大量离散的内存块，同时我们希望这些内存块能按起始地址保证有序，当回收一块内存块时也能尽快恢复有序，那么选用红黑树来存储就很合适了:

set<Mem> freeSet;

构造函数中初始化大小为n的内存块就可以这样表示了：

Allocator(int n) {
    freeSet.insert({0, n});
}

再来看下释放操作，释放时是释放所有id为mID的内存单元，我们也希望释放时能快速找到所有id为mID的内存块,所以可以选用哈希表通过mID指向一个存储了所有该id的内存块的组成的链表:

unordered_map<int, list<Mem>> id2NodeMap;

当分配一块大小为n,id为mID的内存块时，就把该内存块添加到对应链表上。

int allocate(int size, int mID) {
    int insertIdx = -1;
    for(auto it = freeSet.begin(); it != freeSet.end(); it++){
        if((*it).size >= size){
            Mem mem = (*it);
            freeSet.erase(it);
            freeSet.insert({mem.addr + size, mem.size - size});      //旧
            id2NodeMap[mID].push_back({mem.addr, size});             //使用
            insertIdx = mem.addr;
            break;
        }
    }

    return insertIdx;
}

释放时，遍历对应链表，将链表上的内存块依次放回freeSet中，但注意每放回一块内存块时，还需要检查该内存块在freeSet左右是否有相连的内存块，有的话需要合并，好在set提供的迭代器能在O(1)时间找到相邻内存块。

int free(int mID) {
    int num = 0;

    for(auto &mem : id2NodeMap[mID]){
        num += mem.size;
        auto freeIt = freeSet.insert(mem).first;
        auto tmpIt = freeIt;
        bool change = false;
        if(tmpIt!= freeSet.begin()){
            tmpIt--;
            if(tmpIt->addr + tmpIt->size == mem.addr){
                mem.size += tmpIt->size;
                mem.addr = tmpIt->addr;
                freeSet.erase(tmpIt);
                change = true;
            }
        }
        tmpIt = freeIt;
        tmpIt++;
        if(tmpIt != freeSet.end()){
            if(mem.addr + mem.size == tmpIt->addr){
                mem.size += tmpIt->size;
                freeSet.erase(tmpIt);
                change = true;
            }
        }
        if(change){
            freeSet.erase(freeIt);
            freeSet.insert(mem);
        }
    }
    id2NodeMap[mID] = list<Mem>();

    return num;
}

完整代码

struct Mem{
    int addr;
    int size;
};
inline bool operator<(const Mem &m1, const Mem &m2){
    if(m1.addr != m2.addr)
        return m1.addr < m2.addr;
    return m1.size < m2.size;
}
class Allocator {
private:
    set<Mem> freeSet;
    unordered_map<int, list<Mem>> id2NodeMap;
public:
    Allocator(int n) {
        freeSet.insert({0, n});
    }
    
    int allocate(int size, int mID) {
        int insertIdx = -1;
        for(auto it = freeSet.begin(); it != freeSet.end(); it++){
            if((*it).size >= size){
                Mem mem = (*it);
                freeSet.erase(it);
                freeSet.insert({mem.addr + size, mem.size - size});      //旧
                id2NodeMap[mID].push_back({mem.addr, size});             //使用
                insertIdx = mem.addr;
                break;
            }
        }

        return insertIdx;
    }
    
    int free(int mID) {
        int num = 0;

        for(auto &mem : id2NodeMap[mID]){
            num += mem.size;
            auto freeIt = freeSet.insert(mem).first;
            auto tmpIt = freeIt;
            bool change = false;
            if(tmpIt!= freeSet.begin()){
                tmpIt--;
                if(tmpIt->addr + tmpIt->size == mem.addr){
                    mem.size += tmpIt->size;
                    mem.addr = tmpIt->addr;
                    freeSet.erase(tmpIt);
                    change = true;
                }
            }
            tmpIt = freeIt;
            tmpIt++;
            if(tmpIt != freeSet.end()){
                if(mem.addr + mem.size == tmpIt->addr){
                    mem.size += tmpIt->size;
                    freeSet.erase(tmpIt);
                    change = true;
                }
            }
            if(change){
                freeSet.erase(freeIt);
                freeSet.insert(mem);
            }
        }
        id2NodeMap[mID] = list<Mem>();

        return num;
    }
};