testcase-data/get_wrong_status_subset.py

import json
file_path = "/home/luoxianzhen/yang/eval_wrong_code/ALLmode_results/tcb-claude-sonnet-4-20250514-thinking-lcb-fliter-lcb-rank5-all.json"

data = json.load(open(file_path, "r", encoding="utf-8"))

wrong_status = [
    "RE",
    "TLE",
    "MLE",
]
wrong_count = 0
save_code_dict = {}
ti_count = 0
for k, v in data.items():
    save_code_dict[k] = []
    for code in v['codes']:
        save_flag = False
        for status in wrong_status:
            if status in code['status']:
                save_flag = True
        if save_flag:
            save_code_dict[k].append(code['code_id'])
            wrong_count += 1

rest_count = 3000 - wrong_count
for k, v in data.items():
    if k == "似乎在梦中见过的样子":
        a =1
    if len(save_code_dict[k]) <= 0:
        for code in v['codes']:
            save_code_dict[k].append(code['code_id'])
            rest_count -= 1
            if rest_count <= 0:
                break
        if rest_count <= 0:
                break
total = 0
for k, v in save_code_dict.items():
    if len(v) > 0:
        ti_count += 1
        total += len(v)
print(ti_count, total)

print(wrong_count)
json.dump(save_code_dict, open(f"/home/luoxianzhen/yang/eval_wrong_code/cpu_subset.json", "w", encoding="utf-8"), indent=4, ensure_ascii=False)