How to edit a txt file in python-CodePudding

I have a folder containing text files like this:

3 0.39453125 0.49609375 0.017578125 0.0302734375
1 0.4130859375 0.5029296875 0.0224609375 0.025390625
5 0.8125 0.46875 0.1123046875 0.0849609375
0 0.4677734375 0.5078125 0.0283203125 0.0380859375

The first digit is the class. How can I make it so that only lines belonging to class 0 and 1 remain in the files?

So in the end it should look like this:

1 0.4130859375 0.5029296875 0.0224609375 0.025390625
0 0.4677734375 0.5078125 0.0283203125 0.0380859375

I tried this:

for file in os.listdir(folder):
  with open(file,"r ") as f:
      new_file = f.readlines()
      f.seek(0)
      for line in new_file:
          if line.startswith('0') or line.startswith('1'):
              f.write(line)
      f.truncate()

CodePudding user response：

i hope this works for you :)

file = open("text.txt", "r ")

text_list = [line for line in file.readlines() if line.split()[0] in ["0", "1"]]

file.truncate(0)
file.seek(0)
file.write("".join(text_list))
file.close()

CodePudding user response：

The following code should only return line 0 and 1 in an array called write.

write = []   
with open("test.txt",'r',encoding = 'utf-8') as f:
    for line in f:
        if line[0] = '1' or line[0] = '0':
          write  = [line]