LeetCode/rs/meeting-rooms-iii.rs

use std::cmp;
use std::collections::BinaryHeap;

struct MinHeap<T: Ord> {
    heap: BinaryHeap<cmp::Reverse<T>>,
}

impl<T: Ord> MinHeap<T> {
    pub fn new() -> Self {
        Self {
            heap: BinaryHeap::new(),
        }
    }

    pub fn push(&mut self, item: T) {
        self.heap.push(cmp::Reverse(item));
    }

    pub fn pop(&mut self) -> Option<T> {
        self.heap.pop().map(|cmp::Reverse(item)| item)
    }

    pub fn is_empty(&self) -> bool {
        self.heap.is_empty()
    }

    pub fn peek(&self) -> Option<&T> {
        self.heap.peek().map(|cmp::Reverse(item)| item)
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
struct Meeting {
    start: i64,
    end: i64,
}

impl Meeting {
    fn duration(&self) -> i64 {
        self.end - self.start
    }
}

impl From<Vec<i32>> for Meeting {
    fn from(m: Vec<i32>) -> Self {
        Self {
            start: m[0].into(),
            end: m[1].into(),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
struct Allocation {
    end_time: i64,
    room: usize,
}

impl Allocation {
    fn new(end_time: i64, room: usize) -> Self {
        Self { end_time, room }
    }
}

struct Solution {}
impl Solution {
    pub fn most_booked(n: i32, meetings: Vec<Vec<i32>>) -> i32 {
        let mut meetings: Vec<Meeting> = meetings.into_iter().map(|m| m.into()).collect();
        meetings.sort();

        let mut free_rooms: MinHeap<usize> = MinHeap::new();
        (0..n as usize).for_each(|i| free_rooms.push(i));
        let mut nonfree_rooms: MinHeap<Allocation> = MinHeap::new();

        let mut counts: Vec<i32> = vec![0; n as usize];
        for meeting in &meetings {
            // re-add rooms that have been freed
            while let Some(allocation) = nonfree_rooms.peek() {
                if allocation.end_time > meeting.start {
                    break;
                }

                free_rooms.push(allocation.room);
                nonfree_rooms.pop();
            }

            // occupy a room with the meeting
            if let Some(free_room) = free_rooms.pop() {
                nonfree_rooms.push(Allocation::new(meeting.end, free_room));
                counts[free_room] += 1;
            } else if let Some(allocation) = nonfree_rooms.pop() {
                nonfree_rooms.push(Allocation::new(
                    allocation.end_time + meeting.duration(),
                    allocation.room,
                ));
                counts[allocation.room] += 1;
            }
        }

        (0..counts.len()).fold(0, |best_idx, idx| {
            if counts[idx] > counts[best_idx] {
                idx
            } else {
                best_idx
            }
        }) as i32
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn example_1() {
        assert_eq!(
            Solution::most_booked(2, vec![vec![0, 10], vec![1, 5], vec![2, 7], vec![3, 4]]),
            0
        );
    }

    #[test]
    fn example_2() {
        assert_eq!(
            Solution::most_booked(
                3,
                vec![vec![1, 20], vec![2, 10], vec![3, 5], vec![4, 9], vec![6, 8]]
            ),
            1
        );
    }
}