如何用Pandas为文档-词矩阵DataFrame添加文件名列（去后缀）

我通过多个txt文件生成了文档-词矩阵，结果是一个每行对应一个文件、每列对应一个单词的DataFrame（最终目标是用matplotlib可视化该矩阵）。当前DataFrame的索引为数字，但我希望添加一列（或设置索引）为文件名（文件名是年份格式，如“1905.txt”“1906.txt”等），且最好能去除“.txt”后缀。

当前代码如下：

from sklearn.feature_extraction.text import CountVectorizer
from pathlib import Path
import pandas as pd
import numpy as np
import re

# create a list for all txt files
corpus =[]

# with pathlib, get all files in the corpus list 
for fichier in Path("/Users/MyPath/files").rglob("*.txt"):
     corpus.append(fichier.parent / fichier.name)


corpus.sort()
 
all_documents = []
for fichier_txt in corpus:
    with open(fichier_txt) as f:
        fichier_txt_chaine = f.read()
        fichier_txt_chaine = re.sub('[^A-Za-z]', ' ', fichier_txt_chaine) 
    all_documents.append(fichier_txt_chaine)

# here i am using sklearn, but this part is not relevant for my question
coun_vect = CountVectorizer(stop_words= "english")
count_matrix = coun_vect.fit_transform(all_documents)

count_array = count_matrix.toarray()
allDataframe = pd.DataFrame(data=count_array,columns = coun_vect.get_feature_names())
print(allDataframe)
allDataframe.to_csv("Matrice_doc_term.csv")

你可以利用Path对象的stem属性直接获取不带后缀的文件名，在遍历文件时同步收集这些文件名，之后设置为DataFrame的索引或者新增一列。

修改后的代码如下：

from sklearn.feature_extraction.text import CountVectorizer
from pathlib import Path
import pandas as pd
import numpy as np
import re

# create a list for all txt files
corpus =[]

# with pathlib, get all files in the corpus list 
for fichier in Path("/Users/MyPath/files").rglob("*.txt"):
     corpus.append(fichier.parent / fichier.name)


corpus.sort()
 
all_documents = []
# 新增列表存储不带后缀的文件名
file_names = []

for fichier_txt in corpus:
    with open(fichier_txt) as f:
        fichier_txt_chaine = f.read()
        fichier_txt_chaine = re.sub('[^A-Za-z]', ' ', fichier_txt_chaine) 
    all_documents.append(fichier_txt_chaine)
    # 获取不带后缀的文件名，存入列表
    file_names.append(fichier_txt.stem)

# sklearn部分保持不变
coun_vect = CountVectorizer(stop_words= "english")
count_matrix = coun_vect.fit_transform(all_documents)

count_array = count_matrix.toarray()
# 创建DataFrame时直接设置索引为文件名
allDataframe = pd.DataFrame(data=count_array, columns=coun_vect.get_feature_names(), index=file_names)

# 如果想新增一列而不是设置索引，用下面这行替代上面的index参数
# allDataframe['year'] = file_names

print(allDataframe)
allDataframe.to_csv("Matrice_doc_term.csv")

关键说明 #

• fichier_txt.stem：Path对象的stem属性会直接返回文件名（不含后缀），比如1905.txt会返回1905，不需要额外正则处理。
• 同步收集文件名：在遍历读取每个txt文件时，把对应的文件名存入file_names列表，保证顺序和all_documents完全一致，避免错位。
• 两种选择：
  1. 设置为索引：创建DataFrame时指定index=file_names，这样每行的索引就是年份，后续可视化时x轴可以直接用索引。
  2. 新增列：通过allDataframe['year'] = file_names添加一列，适合需要保留原数字索引的场景。

内容的提问来源于stack exchange，提问作者Crazy chicken