Submission #69341816

---

Source Code Expand

import sys

def main():
  try:
    readline = sys.stdin.readline
    
    N_str, R_str = readline().split()
    N = int(N_str)
    R = int(R_str)
    
    L_str_list = readline().split()
    L = [int(x) for x in L_str_list]
  except (IOError, ValueError):
    return

  unlocked_indices = []
  for i in range(N):
    if L[i] == 0:
      unlocked_indices.append(i + 1)

  if not unlocked_indices:
    print(0)
    return

  l = unlocked_indices[0]
  r = unlocked_indices[-1]
  l_room = l - 1

  start_door, end_door = 0, 0
  if R > r:
    start_door = l
    end_door = R
  elif R < l_room:
    start_door = R + 1
    end_door = r
  else:
    start_door = l
    end_door = r

  P = [0] * (N + 1)
  for i in range(N):
    P[i + 1] = P[i] + L[i]

  def count_locked(start, end):
    if start > end:
      return 0
    return P[end] - P[start - 1]

  num_doors_in_range = end_door - start_door + 1
  num_locked_in_range = count_locked(start_door, end_door)

  total_switches = num_doors_in_range + num_locked_in_range
  
  print(total_switches)

if __name__ == "__main__":
  main()

Submission Info

Judge Result