Run isort

library.py

@@ -0,0 +1,126 @@
+import unittest
+from itertools import count
+from math import floor
+
+
+def is_workday(base, offset):
+    return ((base + offset + 1) % 7) > 1
+
+
+def quadratic_larger_root(a, b, c):
+    """Find the larger root of the quadratic equation"""
+    b /= a
+    c /= a
+    # (x + b/2) ** 2 = -c + b *b / 4
+    r = (-c + b * b / 4) ** 0.5
+    return -b / 2 + r
+
+
+def bsearch_solve(f):
+    """Solve the equation f(x) == 0 assuming f(0) > 0 and the function decreases"""
+    l = 0
+    r = 1
+    while f(r) >= 0:
+        r *= 3
+
+    while r - l > 1:
+        assert f(l) >= 0 > f(r)
+        m = (l + r) // 2
+        if f(m) >= 0:
+            l = m
+        else:
+            r = m
+
+    return l
+
+
+class LibrarySolver:
+    def __init__(self, use_quadratic=True):
+        self.use_quadratic = use_quadratic
+
+    def solve(self, k, m, d) -> int:
+        """Solve the library problem using one of the two methods"""
+
+        def simulate(from_books, from_week, to_week=None):
+            left_books = from_books
+            start_day = 7 * from_week + 1
+
+            if to_week is None:
+                days = count(start_day)
+            else:
+                days = range(start_day, 7 * to_week + 1)
+
+            for day in days:
+                left_books += (k if is_workday(d, day - 1) else 0) - day
+                if left_books < 0:
+                    return day - 1
+
+        first_week = simulate(m, 0, 1)
+        if first_week is not None:
+            return first_week
+
+        # f(x) = needs - has
+        # f(0) <= 0 and f(infty) > 0
+        # arg min f between 0 and 1
+        k_needs = 49, 7, 0
+        k_has = 0, 10 * k, 2 * m
+        k_coef = [n - h for n, h in zip(k_needs, k_has)]
+
+        def left_after_weeks(weeks):
+            return -(k_coef[0] * weeks**2 + k_coef[1] * weeks + k_coef[2]) // 2
+
+        if self.use_quadratic:
+            weeks = floor(quadratic_larger_root(*k_coef))
+        else:
+            weeks = bsearch_solve(left_after_weeks)
+
+        left = left_after_weeks(weeks)
+        assert left >= 0
+
+        return simulate(left, weeks)
+
+
+class TestLibrarySolution(unittest.TestCase):
+    def test_workday(self):
+        self.assertFalse(is_workday(6, 1))
+        self.assertTrue(is_workday(6, 2))
+        self.assertTrue(is_workday(6, 6))
+        self.assertFalse(is_workday(6, 7))
+        self.assertTrue(is_workday(1, 0))
+        self.assertFalse(is_workday(6, 0))
+        self.assertFalse(is_workday(7, 0))
+        self.assertFalse(is_workday(6, 1))
+        self.assertTrue(is_workday(7, 1))
+
+    def test_quadratic_larger_root(self):
+        self.assertEqual(quadratic_larger_root(1, 0, -1), 1)
+        self.assertEqual(quadratic_larger_root(1, -2, 1), 1)
+
+    def test_solve_library(self):
+        for parameter in [False, True]:
+            solve_library = LibrarySolver(parameter).solve
+            self.assertEqual(solve_library(1, 1, 1), 2)
+            self.assertEqual(solve_library(1, 1, 2), 2)
+            self.assertEqual(solve_library(1, 1, 3), 2)
+            self.assertEqual(solve_library(1, 1, 4), 2)
+            self.assertEqual(solve_library(1, 1, 5), 1)
+            self.assertEqual(solve_library(1, 1, 6), 1)
+            self.assertEqual(solve_library(1, 1, 7), 1)
+            self.assertEqual(solve_library(1, 2, 7), 2)
+            self.assertEqual(solve_library(2, 1, 7), 2)
+            self.assertEqual(solve_library(2, 1, 2), 3)
+
+    def test_solve_fails_immediately(self):
+        for parameter in [False, True]:
+            solve_library = LibrarySolver(parameter).solve
+            self.assertEqual(solve_library(0, 0, 6), 0)
+            self.assertEqual(solve_library(1, 0, 6), 0)
+            self.assertEqual(solve_library(100, 0, 6), 0)
+
+
+def main():
+    print(LibrarySolver().solve(*(map(int, input().split()))))
+
+
+if __name__ == "__main__":
+    main()

trains.py

@@ -0,0 +1,118 @@
+import unittest
+
+
+class TestFenwickTree(unittest.TestCase):
+    def setUp(self):
+        self.t = FenwickTree(3)
+
+    def test_add_and_upto(self):
+        self.t.add(2)
+        self.assertEqual(self.t.upto(0), 0)
+        self.assertEqual(self.t.upto(1), 0)
+        self.assertEqual(self.t.upto(2), 1)
+        self.assertEqual(self.t.upto(3), 1)
+
+    def test_min_free(self):
+        self.t.add(2)
+        self.assertEqual(self.t.min_free(), 1)
+        self.t.add(1)
+        self.t.add(3)
+        self.assertEqual(self.t.min_free(), 4)
+
+    def test_remove(self):
+        self.t.add(2)
+        self.t.add(1)
+        self.assertEqual(self.t.min_free(), 3)
+        self.t.add(3)
+        self.t.remove(2)
+        self.assertEqual(self.t.min_free(), 2)
+
+
+class FenwickTree:
+    """Simple Fenwick tree. The keys are numbered from 1 to 2^k."""
+
+    def __init__(self, N):
+        while N != (N & -N):
+            N += N & -N
+        self.N = N
+        self.cell = [0] * (N + 1)
+
+    def add(self, key, value=1):
+        """Add a value to the specified key in the tree."""
+        while key <= self.N:
+            self.cell[key] += value
+            key += key & -key
+
+    def remove(self, key):
+        """Subtract 1 from the value of the specified key in the tree."""
+        self.add(key, -1)
+
+    def upto(self, key):
+        """Sum the values of the keys that are smaller or equal to the given tree."""
+        s = 0
+        c = min(key, self.N)
+        while c:
+            s += self.cell[c]
+            c -= c & -c
+        return s
+
+    def min_free(self):
+        """Find the smallest key for which the value is 0.
+
+        Requirement: all values are 0 or 1.
+
+        :return: The minimum free key (self.N + 1 if all keys are set).
+        """
+        r = 1
+        while self.cell[r] == r:
+            r *= 2
+            if r > self.N:
+                return self.N + 1
+
+        b = r // 2
+        while b > 1:
+            # [r - b + 1; r] has the answer
+            b = b // 2
+            m = r - b
+            if self.cell[m] < b:
+                r = m
+            else:
+                assert self.cell[m] == b
+
+        return r
+
+
+def main():
+    k, n = map(int, input().split())
+
+    def gen():
+        for train in range(n):
+            arrive, leave = [int(s) for s in input().split()]
+            yield arrive, False, train
+            yield leave, True, train
+
+    events = sorted(list(gen()))
+
+    # platforms are numbered from 1 to k
+    platforms = FenwickTree(k)
+    # trains are numbered from 0 to n-1
+    trains = []
+
+    def find_platforms():
+        for _, leave, train in events:
+            if leave:
+                assert train < len(trains)
+                platforms.remove(trains[train])
+            elif (platform := platforms.min_free()) <= k:
+                assert len(trains) == train
+                trains.append(platform)
+                platforms.add(platform)
+            else:
+                return 0, train + 1
+        return trains
+
+    print(*find_platforms())
+
+
+if __name__ == "__main__":
+    main()