import sys import numpy def main(): data = sys.stdin.read().split() if not data: return n = int(data[0]) k = int(data[1]) t = int(data[2]) s = data[3].strip() s_arr = numpy.array(list(s)) free = (s_arr == '3').astype(numpy.int64) t1 = (s_arr == '1').astype(numpy.int64) t2 = (s_arr == '2').astype(numpy.int64) meeting = ((s_arr == '1') | (s_arr == '2')).astype(numpy.int64) Pf = numpy.cumsum(free) Pt1 = numpy.cumsum(t1) Pt2 = numpy.cumsum(t2) Pm = numpy.cumsum(meeting) total_free = int(Pf[-1]) total_t1 = int(Pt1[-1]) total_t2 = int(Pt2[-1]) if total_t1 > k: no_trip = -10**9 else: budget_no_trip = k - total_t1 no_trip = total_free + total_t1 + min(total_t2, budget_no_trip) best = no_trip LF = numpy.zeros(n+1, dtype=numpy.int64) LT1 = numpy.zeros(n+1, dtype=numpy.int64) LT2 = numpy.zeros(n+1, dtype=numpy.int64) LF[1:] = Pf[:-1] LT1[1:] = Pt1[:-1] LT2[1:] = Pt2[:-1] RF = numpy.zeros(n, dtype=numpy.int64) RT1 = numpy.zeros(n, dtype=numpy.int64) RT2 = numpy.zeros(n, dtype=numpy.int64) RF[:-1] = total_free - Pf[:-1] RF[-1] = 0 RT1[:-1] = total_t1 - Pt1[:-1] RT1[-1] = 0 RT2[:-1] = total_t2 - Pt2[:-1] RT2[-1] = 0 maxL = n - t T_left = numpy.empty(maxL+1, dtype=numpy.int64) T_left[0] = Pm[t-1] if maxL >= 1: T_left[1:] = Pm[t: n] - Pm[:maxL] for R in range(2*t - 1, n): L_upper = min(maxL, R - (2*t - 1)) if L_upper < 0: continue cand_L = numpy.arange(0, L_upper+1) travel_left = T_left[cand_L] travel_right = int(Pm[R] - (Pm[R-t] if R-t >= 0 else 0)) travel_cost = travel_left + travel_right B = k - travel_cost left_free = LF[cand_L] left_t1 = LT1[cand_L] left_t2 = LT2[cand_L] right_free = int(RF[R]) right_t1 = int(RT1[R]) right_t2 = int(RT2[R]) home_free = left_free + right_free home_t1 = left_t1 + right_t1 home_t2 = left_t2 + right_t2 feasible = B >= home_t1 if not numpy.any(feasible): continue extra = numpy.minimum(home_t2, numpy.maximum(0, B - home_t1)) home_points = home_free + home_t1 + extra candidate_points = home_points[feasible] if candidate_points.size > 0: best_candidate = candidate_points.max() if best_candidate > best: best = best_candidate if best < 0: best = -1 print(best) if __name__ == '__main__': main()
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | import sys import numpy def main(): data = sys.stdin.read().split() if not data: return n = int(data[0]) k = int(data[1]) t = int(data[2]) s = data[3].strip() s_arr = numpy.array(list(s)) free = (s_arr == '3').astype(numpy.int64) t1 = (s_arr == '1').astype(numpy.int64) t2 = (s_arr == '2').astype(numpy.int64) meeting = ((s_arr == '1') | (s_arr == '2')).astype(numpy.int64) Pf = numpy.cumsum(free) Pt1 = numpy.cumsum(t1) Pt2 = numpy.cumsum(t2) Pm = numpy.cumsum(meeting) total_free = int(Pf[-1]) total_t1 = int(Pt1[-1]) total_t2 = int(Pt2[-1]) if total_t1 > k: no_trip = -10**9 else: budget_no_trip = k - total_t1 no_trip = total_free + total_t1 + min(total_t2, budget_no_trip) best = no_trip LF = numpy.zeros(n+1, dtype=numpy.int64) LT1 = numpy.zeros(n+1, dtype=numpy.int64) LT2 = numpy.zeros(n+1, dtype=numpy.int64) LF[1:] = Pf[:-1] LT1[1:] = Pt1[:-1] LT2[1:] = Pt2[:-1] RF = numpy.zeros(n, dtype=numpy.int64) RT1 = numpy.zeros(n, dtype=numpy.int64) RT2 = numpy.zeros(n, dtype=numpy.int64) RF[:-1] = total_free - Pf[:-1] RF[-1] = 0 RT1[:-1] = total_t1 - Pt1[:-1] RT1[-1] = 0 RT2[:-1] = total_t2 - Pt2[:-1] RT2[-1] = 0 maxL = n - t T_left = numpy.empty(maxL+1, dtype=numpy.int64) T_left[0] = Pm[t-1] if maxL >= 1: T_left[1:] = Pm[t: n] - Pm[:maxL] for R in range(2*t - 1, n): L_upper = min(maxL, R - (2*t - 1)) if L_upper < 0: continue cand_L = numpy.arange(0, L_upper+1) travel_left = T_left[cand_L] travel_right = int(Pm[R] - (Pm[R-t] if R-t >= 0 else 0)) travel_cost = travel_left + travel_right B = k - travel_cost left_free = LF[cand_L] left_t1 = LT1[cand_L] left_t2 = LT2[cand_L] right_free = int(RF[R]) right_t1 = int(RT1[R]) right_t2 = int(RT2[R]) home_free = left_free + right_free home_t1 = left_t1 + right_t1 home_t2 = left_t2 + right_t2 feasible = B >= home_t1 if not numpy.any(feasible): continue extra = numpy.minimum(home_t2, numpy.maximum(0, B - home_t1)) home_points = home_free + home_t1 + extra candidate_points = home_points[feasible] if candidate_points.size > 0: best_candidate = candidate_points.max() if best_candidate > best: best = best_candidate if best < 0: best = -1 print(best) if __name__ == '__main__': main() |