protohack

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs

mean.rs (2763B)


      1 use std::collections::BTreeMap;
      2 use tokio::io::{AsyncReadExt, AsyncWriteExt, stdin};
      3 use tokio::sync::oneshot;
      4 
      5 #[tokio::main]
      6 async fn main() -> std::io::Result<()> {
      7     let (tx, rx) = oneshot::channel();
      8 
      9     tokio::spawn(async move {
     10         let mut reader = stdin();
     11         let mut buffer = [0u8; 1];
     12         loop {
     13             if reader.read_exact(&mut buffer).await.is_ok() {
     14                 if buffer[0] == b'q' {
     15                     let _ = tx.send(());
     16                     break;
     17                 }
     18             }
     19         }
     20     });
     21 
     22     let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
     23     println!("Server started. Press 'q' then Enter to stop.");
     24 
     25     tokio::select! {
     26         res = async {
     27             loop {
     28                 let (mut socket, _) = listener.accept().await?;
     29                 tokio::spawn(async move {
     30                     let (mut rd, mut wr) = socket.split();
     31                     let mut buf = [0u8; 9];
     32                     let mut history = BTreeMap::new();
     33                     loop {
     34                         if rd.read_exact(&mut buf).await.is_err() {
     35                             break;
     36                         }
     37                         match buf[0] {
     38                             b'I' => {
     39                                 let t = i32::from_be_bytes(buf[1..5].try_into().unwrap());
     40                                 let p = i32::from_be_bytes(buf[5..9].try_into().unwrap());
     41                                 history.insert(t, p);
     42                             }
     43                             b'Q' => {
     44                                 let min = i32::from_be_bytes(buf[1..5].try_into().unwrap());
     45                                 let max = i32::from_be_bytes(buf[5..9].try_into().unwrap());
     46                                 let avg = calculate_average(&history, min, max);
     47                                 if wr.write_all(&avg.to_be_bytes()).await.is_err() {
     48                                     break;
     49                                 }
     50                             }
     51                             _ => break,
     52                         }
     53                     }
     54                 });
     55             }
     56             #[allow(unreachable_code)]
     57             Ok::<(), std::io::Error>(())
     58         } => {
     59             if let Err(e) = res {
     60                 eprintln!("Accept error: {}", e);
     61             }
     62         },
     63         _ = rx => {
     64             println!("Shutdown signal received. Exiting...");
     65         }
     66     }
     67 
     68     Ok(())
     69 }
     70 
     71 fn calculate_average(history: &BTreeMap<i32, i32>, min: i32, max: i32) -> i32 {
     72     if min > max {
     73         return 0;
     74     }
     75     let subset: Vec<_> = history.range(min..=max).map(|(_, &v)| v as i64).collect();
     76     if subset.is_empty() {
     77         return 0;
     78     }
     79     (subset.iter().sum::<i64>() / subset.len() as i64) as i32
     80 }