python 成语接龙

# coding=utf-8
import xpinyin
import random
import pickle
from collections import defaultdict
#读文件获取词语列表
file = open('/tensorflow/py_aiplat_demo/data/ciyu.txt','r',encoding='gbk')
line = file.readline()
#定义词语数组
wordList=[]
while line:
    singledata = ''.join(line.split()).split('/')[0]
    wordList.append(singledata)
    line = file.readline()
p=xpinyin.Pinyin()
#定义一个字典类型
poem_dict = defaultdict(list)
#对语料处理
for word in wordList:
    head = p.get_pinyin(word, tone_marks='marks', splitter=' ').split()[0]
    poem_dict[head].append(word)
#print(poem_dict)
#成语接龙开始
selected_words = []
start_word = '四'
start_pinyin = p.get_pinyin(start_word, tone_marks='marks', splitter=' ').split()[0]
for i in range(50):
    #循环处理第一个词语
    start_data=[poem_dict[start_pinyin] for w in poem_dict.keys() if w == start_pinyin]
    if len(start_data)==0:
       break
    #取列表中的数组值
    current_word=start_data[random.randint(0,len(start_data)-1)]
    #去重
    work_word=''
    for word in current_word:
        if word not in selected_words:
            work_word=word
    #接龙中的词的最后一个字转成拼音
    start_pinyin = p.get_pinyin(work_word, tone_marks='marks', splitter=' ').split()[3]
    #将接龙的词存入返回的列表中
    selected_words.append(work_word)
#打印也接龙词列表
print(selected_words)