leetcode 1912. 设计电影租借系统

题目：1912. 设计电影租借系统 - 力扣（LeetCode）

维护有序队列，因为每次查询最多有5个结果，用最小堆就可以，每次查找时bfs一下最小堆。

struct Node {
    int shop;
    int movie;
    int price;
    size_t idx;
    Node(int shop, int movie, int price) {
        this->shop = shop;
        this->movie = movie;
        this->price = price;
    }
};
class MovieRentingSystem {
public:
    MovieRentingSystem(int n, vector<vector<int>>& entries) {
        shopEntries.resize(n);
        for (int i = 0; i < entries.size(); i++) {
            auto& arr = entries[i];
            Node* node = new Node(arr[0], arr[1], arr[2]);
            auto& heap = movieHeaps[node->movie];
            node->idx = heap.size();
            heap.push_back(node);
            HeapUp(heap, node->idx);
            
            (shopEntries[node->shop])[node->movie] = node;
        }
    }
    
    vector<int> search(int movie) {
        vector<int> ret;
        vector<Node*> tmp;
        
        if (movieHeaps.find(movie) == movieHeaps.end()) {
            return ret;
        }
        vector<Node*>& heap = movieHeaps[movie];
        if (heap.empty()) {
            return ret;
        }
        tmp.push_back(heap[0]);
        for (int i = 0; i < 5; i++) {
            if (tmp.empty()) {
                return ret;
            }
            Node* min = tmp[0];
            int minj = 0;
            for (int j = 1; j < tmp.size(); j++) {
                if (AisSmallerThanB(tmp[j], min)) {
                    min = tmp[j];
                    minj = j;
                }
            }
            ret.push_back(min->shop);
            tmp[minj] = tmp[tmp.size() - 1];
            tmp.pop_back();
            if (min->idx * 2 + 1 < heap.size()) {
                tmp.push_back(heap[min->idx * 2 + 1]);
            }
            if (min->idx * 2 + 2 < heap.size()) {
                tmp.push_back(heap[min->idx * 2 + 2]);
            }
        }
        
        return ret;
    }
    
    void rent(int shop, int movie) {
        Node* node = (shopEntries[shop])[movie];
        vector<Node*>& heap = movieHeaps[movie];
        if (heap.size() == 1) {
            heap.clear();
        } else {
            Node* tail = heap[heap.size() - 1];
            heap[node->idx] = tail;
            tail->idx = node->idx;
            heap.pop_back();
            HeapDown(heap, tail->idx);
            HeapUp(heap, tail->idx);
        }
        node->idx = rentHeaps.size();
        rentHeaps.push_back(node);
        HeapUp(rentHeaps, node->idx);
    }
    
    void drop(int shop, int movie) {
        Node* node = (shopEntries[shop])[movie];
        vector<Node*>& heap = movieHeaps[movie];
        if (rentHeaps.size() == 1) {
            rentHeaps.clear();
        } else {
            Node* tail = rentHeaps[rentHeaps.size() - 1];
            rentHeaps[node->idx] = tail;
            tail->idx = node->idx;
            rentHeaps.pop_back();
            HeapDown(rentHeaps, tail->idx);
            HeapUp(rentHeaps, tail->idx);
        }
        node->idx = heap.size();
        heap.push_back(node);
        HeapUp(heap, node->idx);
    }
    
    vector<vector<int>> report() {
        vector<vector<int>> ret;
        vector<Node*> tmp;
        
        vector<Node*>& heap = rentHeaps;
        if (heap.empty()) {
            return ret;
        }
        tmp.push_back(heap[0]);
        for (int i = 0; i < 5; i++) {
            if (tmp.empty()) {
                return ret;
            }
            Node* min = tmp[0];
            int minj = 0;
            for (int j = 1; j < tmp.size(); j++) {
                if (AisSmallerThanB(tmp[j], min)) {
                    min = tmp[j];
                    minj = j;
                }
            }
            vector<int> a;
            a.push_back(min->shop);
            a.push_back(min->movie);
            ret.push_back(a);
            tmp[minj] = tmp[tmp.size() - 1];
            tmp.pop_back();
            if (min->idx * 2 + 1 < heap.size()) {
                tmp.push_back(heap[min->idx * 2 + 1]);
            }
            if (min->idx * 2 + 2 < heap.size()) {
                tmp.push_back(heap[min->idx * 2 + 2]);
            }
        }
        
        return ret;
    }
    
private:
    bool AisSmallerThanB(Node* a, Node* b) {
        if (a->price < b->price) {
            return true;
        }
        if (a->price == b->price && a->shop < b->shop) {
            return true;
        }
        if (a->price == b->price && a->shop == b->shop && a->movie < b->movie) {
            return true;
        }
        return false;
    }
    
    void HeapUp(vector<Node*>& heap, size_t idx) {
        if (idx == 0) {
            return;
        }
        size_t p = (idx - 1) / 2;
        Node* thiz = heap[idx];
        Node* parent = heap[p];
        if (AisSmallerThanB(thiz, parent)) {
            heap[idx] = parent;
            parent->idx = idx;
            heap[p] = thiz;
            thiz->idx = p;
            HeapUp(heap, p);
        }
    }
    
    void HeapDown(vector<Node*>& heap, size_t idx) {
        Node* thiz = heap[idx];
        Node* left = nullptr;
        Node* right = nullptr;
        if (idx * 2 + 1 < heap.size()) {
            left = heap[idx * 2 + 1];
        }
        if (idx * 2 + 2 < heap.size()) {
            right = heap[idx * 2 + 2];
        }
        Node* child = left;
        if (right && AisSmallerThanB(right, left)) {
            child = right;
        }
        if (child && AisSmallerThanB(child, thiz)) {
            size_t c = child->idx;
            heap[idx] = child;
            child->idx = idx;
            heap[c] = thiz;
            thiz->idx = c;
            HeapDown(heap, c);
        }
    }
    
private:
    // moive, heap
    map<int, vector<Node*>> movieHeaps;
    
    vector<Node*> rentHeaps;
    
    vector<map<int, Node*>> shopEntries;
};