ABC kyuridenamida#185

atcodertools/fmtprediction/tokenize_format.py

@@ -70,14 +70,13 @@ def _sanitized_tokens(input_format: str) -> List[str]:
 #    print("PRE-FILTERED TOKENS: " + input_format)
     a = []
     for s in input_format.split('\n'):
-        # AGC007-A, ABC173-C
-        result = re.match('(.)_\{([^,]*)[,]*1\}[^\s].*_\{.*\}', s)
-        if result is not None:
-            s = '{}_{{{}}}'.format(result.group(1), result.group(2))
-        # ARC005-B, ARC005-C
-        result = re.match('(.)_\{\(([^,]*)[,]*[01]\)\}[^\s].*_\{.*\}', s)
-        if result is not None:
-            s = '{}_{{{}}}'.format(result.group(1), result.group(2))
+        # AGC007 A: A_{11}A_{12}...A_{1W} -> A_1
+        # ABC173 C: c_{1,1}c_{1,2}...c_{1,W} -> c_1
+        s = re.sub(r'(.)_{([^,]*)[,]?1}[^\s].*_{[,\w]*}', r'\1_\2', s)
+        # ARC005 C: c_{(0,0)}c_{(0,1)} … c_{(0,W-1)} -> c_0
+        s = re.sub(r'(.)_{\(([^,]*)[,]?[01]\)}.*_{[^{}]*}', r'\1_\2', s)
+        # ABC185 A: 
+        s = re.sub(r'\\hspace{\d+pt}', r'', s)
         a.append(s)
     input_format = '\n'.join(a)
 #    print("FILTERED TOKENS: " + input_format)