protohackers

My solutions to the protohackers.com challenges.
git clone git://henryandlizzy.uk/protohackers
Log | Files | Refs

protohack-2.cpp (2492B)

      1 #include "inet.hpp"
      2 
      3 #include <poll.h>
      4 
      5 #include <algorithm>
      6 #include <iostream>
      7 #include <map>
      8 #include <vector>
      9 
     10 using namespace std;
     11 
     12 struct listings : map<int32_t, int32_t>
     13 {
     14 	int32_t average(int32_t min, int32_t max)
     15 	{
     16 		if (empty())
     17 			return 0;
     18 		if (min > max)
     19 			return 0;
     20 
     21 		auto it = lower_bound(min);
     22 
     23 		if (it == end())
     24 			return 0;
     25 
     26 		int64_t accum = 0;
     27 		int32_t n = 0;
     28 
     29 		while (it != end() and it->first <= max)
     30 		{
     31 			++n;
     32 			accum += it++->second;
     33 		}
     34 		if (n == 0)
     35 			return 0;
     36 
     37 		return accum / n;
     38 	}
     39 };
     40 
     41 struct client
     42 {
     43 	descriptor d;
     44 	listings l;
     45 	std::vector<char> buf;
     46 };
     47 
     48 std::map<int, client> clients;
     49 
     50 uint32_t decode(void const* in)
     51 {
     52 	auto p = (unsigned char const*)in;
     53 	return p[3] | p[2] << 8U | p[1] << 16U | p[0] << 24U;
     54 }
     55 
     56 void encode(uint32_t n, void* out)
     57 {
     58 	unsigned char* p = (unsigned char*)out;
     59 	p[0] = (n & 0xFF000000U) >> 24U;
     60 	p[1] = (n & 0x00FF0000U) >> 16U;
     61 	p[2] = (n & 0x0000FF00U) >> 8U;
     62 	p[3] = (n & 0x000000FFU);
     63 }
     64 
     65 int main()
     66 {
     67 	auto incoming = inet::listen(SOCK_STREAM);
     68 
     69 	for (;;)
     70 	{
     71 		std::vector<pollfd> awaitables;
     72 		awaitables.push_back(pollfd{incoming, POLLIN, 0});
     73 		for (auto const& c : clients)
     74 			awaitables.push_back(pollfd{c.second.d, POLLIN, 0});
     75 
     76 		::poll(awaitables.data(), awaitables.size(), INFTIM);
     77 
     78 		for (auto& ev : awaitables)
     79 		{
     80 			if (~ev.revents & POLLIN)
     81 				continue;
     82 
     83 			if (ev.fd == incoming)
     84 			{
     85 				auto conn = inet::accept(incoming);
     86 				clients[conn].d = std::move(conn);
     87 				continue;
     88 			}
     89 
     90 			if (auto it = clients.find(ev.fd); it != clients.end())
     91 			{
     92 				char buf;
     93 				int rc = ::read(ev.fd, &buf, sizeof buf);
     94 				if (not rc)
     95 				{
     96 					std::cout << "* " << ev.fd << " disconnected\n";
     97 					clients.erase(it);
     98 					continue;
     99 				}
    100 				if (rc < 0)
    101 					err(1, "read");
    102 
    103 				auto& buffer = it->second.buf;
    104 				buffer.push_back(buf);
    105 				if (buffer.size() < 9)
    106 					continue;
    107 
    108 				auto& l = it->second.l;
    109 
    110 				int32_t a = decode(buffer.data() + 1);
    111 				int32_t b = decode(buffer.data() + 5);
    112 				char cmd = buffer.front();
    113 				buffer.erase(buffer.begin(), buffer.begin() + 9);
    114 				switch (cmd)
    115 				{
    116 				case 'I':
    117 					if (l.count(a))
    118 						return 1;
    119 					l[a] = b;
    120 					continue;
    121 
    122 				case 'Q':
    123 					a = l.average(a, b);
    124 					char txbuf[4];
    125 					encode(a, txbuf);
    126 					::write(ev.fd, txbuf, 4);
    127 					continue;
    128 				}
    129 
    130 				std::cout << "* " << ev.fd << " bad command" << cmd << "\n";
    131 				clients.erase(it);
    132 				continue;
    133 			}
    134 
    135 			std::cerr << "Unknown fd " << ev.fd << std::endl;
    136 			return 1;
    137 		}
    138 	}
    139 }