Getting input and output data

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
import os

contest_id, problem_letter, test_case = input().split(' '); test_case = int(test_case)

option = webdriver.ChromeOptions()
option.add_argument('headless')

PATH = r"C:\Users\Vito\Programs\chromedriver.exe"
driver = webdriver.Chrome(PATH, options=option)

driver.get(f"https://www.codeforces.com/contest/{contest_id}/problem/{problem_letter}")

os.system("cmd /c cls")

problem = driver.find_element(by=By.CLASS_NAME, value="title")

print("\nCollected input and output data for problem:", problem.text, sep="\n")

input_file = open("input.txt", "w")
correct_output_file = open("correct_output.txt", "w")

io = driver.find_elements(by=By.TAG_NAME, value="pre")

input_file.write(io[test_case*2-2].text+"\n")
correct_output_file.write(io[test_case*2-1].text+"\n")
input_file.close()
correct_output_file.close()

driver.quit()