def solve(): import sys, numpy as np data = sys.stdin.read().split() if not data: return it = iter(data) n = int(next(it)); k = int(next(it)); t = int(next(it)) obligations = next(it).strip() f = np.array([1 if ch == '1' else 0 for ch in obligations], dtype=np.int32) o = np.array([1 if ch == '2' else 0 for ch in obligations], dtype=np.int32) g = np.array([1 if ch == '3' else 0 for ch in obligations], dtype=np.int32) T = np.array([1 if ch in '12' else 0 for ch in obligations], dtype=np.int32) pf = np.empty(n + 1, dtype=np.int32) po = np.empty(n + 1, dtype=np.int32) pg = np.empty(n + 1, dtype=np.int32) pT = np.empty(n + 1, dtype=np.int32) pf[0] = po[0] = pg[0] = pT[0] = 0 for i in range(n): pf[i + 1] = pf[i] + f[i] po[i + 1] = po[i] + o[i] pg[i + 1] = pg[i] + g[i] pT[i + 1] = pT[i] + T[i] total_f, total_o, total_g = pf[n], po[n], pg[n] no_trip = -10**9 if total_f <= k: no_trip = total_g + total_f + min(total_o, k - total_f) best_trip = -10**9 max_s = n - 2 * t for s in range(max_s + 1): r_start = s + 2 * t - 1 if r_start >= n: continue r_arr = np.arange(r_start, n, dtype=np.int32) transit_out = pT[s + t] - pT[s] transit_in = pT[r_arr + 1] - pT[r_arr - t + 1] transit_cost = transit_out + transit_in home_f = total_f - (pf[r_arr + 1] - pf[s]) home_o = total_o - (po[r_arr + 1] - po[s]) home_g = total_g - (pg[r_arr + 1] - pg[s]) avail = k - transit_cost feasible = avail >= home_f if not np.any(feasible): continue candidate = home_g[feasible] + home_f[feasible] + np.minimum(home_o[feasible], avail[feasible] - home_f[feasible]) best_trip = max(best_trip, candidate.max()) ans = max(no_trip, best_trip) if ans < 0: ans = -1 sys.stdout.write(str(int(ans))) if __name__ == '__main__': solve()
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 | def solve(): import sys, numpy as np data = sys.stdin.read().split() if not data: return it = iter(data) n = int(next(it)); k = int(next(it)); t = int(next(it)) obligations = next(it).strip() f = np.array([1 if ch == '1' else 0 for ch in obligations], dtype=np.int32) o = np.array([1 if ch == '2' else 0 for ch in obligations], dtype=np.int32) g = np.array([1 if ch == '3' else 0 for ch in obligations], dtype=np.int32) T = np.array([1 if ch in '12' else 0 for ch in obligations], dtype=np.int32) pf = np.empty(n + 1, dtype=np.int32) po = np.empty(n + 1, dtype=np.int32) pg = np.empty(n + 1, dtype=np.int32) pT = np.empty(n + 1, dtype=np.int32) pf[0] = po[0] = pg[0] = pT[0] = 0 for i in range(n): pf[i + 1] = pf[i] + f[i] po[i + 1] = po[i] + o[i] pg[i + 1] = pg[i] + g[i] pT[i + 1] = pT[i] + T[i] total_f, total_o, total_g = pf[n], po[n], pg[n] no_trip = -10**9 if total_f <= k: no_trip = total_g + total_f + min(total_o, k - total_f) best_trip = -10**9 max_s = n - 2 * t for s in range(max_s + 1): r_start = s + 2 * t - 1 if r_start >= n: continue r_arr = np.arange(r_start, n, dtype=np.int32) transit_out = pT[s + t] - pT[s] transit_in = pT[r_arr + 1] - pT[r_arr - t + 1] transit_cost = transit_out + transit_in home_f = total_f - (pf[r_arr + 1] - pf[s]) home_o = total_o - (po[r_arr + 1] - po[s]) home_g = total_g - (pg[r_arr + 1] - pg[s]) avail = k - transit_cost feasible = avail >= home_f if not np.any(feasible): continue candidate = home_g[feasible] + home_f[feasible] + np.minimum(home_o[feasible], avail[feasible] - home_f[feasible]) best_trip = max(best_trip, candidate.max()) ans = max(no_trip, best_trip) if ans < 0: ans = -1 sys.stdout.write(str(int(ans))) if __name__ == '__main__': solve() |