diff --git a/library.py b/library.py new file mode 100644 index 0000000..6c5212f --- /dev/null +++ b/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() diff --git a/trains.py b/trains.py new file mode 100644 index 0000000..d6d469c --- /dev/null +++ b/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()