去重并保留第n个重复元素（python）

题目

看到一道面试题：题目出处

数组去重，对每个重复的元素，保留第n个重复位置上的元素，其他的都删除
若元素的重复次数小于n，则保留最后一个重复位置上的元素

例如数组[ 2, 2, 2, 3, 3, 2, 2, 4, 5, 3, 4, 7, 8, 1 ]，n=3
要求输出为[ 2, 5, 3, 4, 7, 8, 1 ]

思路

总体思路：遍历一次记录数组中出现的元素，以及它们出现的位置；再从中取得要保留的位置的元素，生成一个新数组作为最终结果

设置一个列表idxs，存放多个子列表，每个子列表是某个元素在数组中所有出现的位置
设置一个哈希表hashDict，存放数组中出现过的元素，键值为元素，值为idxs中存放这个元素出现位置的子列表的下标
图示：

遍历1次数组arr，生成hashDict和idxs：

• 若当前元素arr[i]不在hashDict中，在hashDict中追加{arr[i]:len(idxs)}，并向idxs中追加一个新列表[i]
  表示：arr[i]这个元素的出现位置存放在idxs中要新增的一个列表，这个列表现在只有一个值i，表示arr[i]在i位置出现

• 若当前元素arr[i]已在hashDict中，查hashDict[arr[i]]，可以得到它的出现位置应该存放到idxs的哪个子列表里，将i追加到idxs[hashDict[arr[i]]]中

得到hashDict和idxs后，设置1个新列表dst，存放之后要保留的元素的下标

遍历1次hashDict，检查每个元素出现过多少次：

• ≥3次，把该元素出现的第三个位置下标加入dst
• <3次，把该元素出现的最后一个位置下标加入dst

对dst排序，使其中的下标递增，再遍历一遍dst，把每个元素替换成对应的arr中的元素，即可得到最终结果

---

复杂度分析

用了排序，时间复杂度O(nlgn)
用了几个辅助数组记录索引，空间复杂度O(n)

没找到对应的原题，只是个人解法，仅供参考，不一定是最优

代码

def saveN(arr, n):
    # 去重，并且对每个重复元素，保留其第n个重复位置上的元素，其它位置上该元素删除
    # 若n大于重复元素的个数，则保留最后一个重复元素
    
    # dst用来存放最终的去重结果，
    # idxs存放多个子列表，每个列表是一个重复元素的所有出现位置
    dst, idxs = [], []
    # hashDict存放数组中出现过的元素，和idxs中存放他们下标出现位置的子列表的下标
    hashDict = {}
    for i in range(len(arr)):
        if arr[i] not in hashDict:
            hashDict.update({arr[i]: len(idxs)})
            idxs.append([i])
        else:
            idxs[hashDict[arr[i]]].append(i)
    for i in hashDict:
        if len(idxs[hashDict[i]]) >= n:
            dst.append(idxs[hashDict[i]][2])
        else:
            dst.append(idxs[hashDict[i]][-1])
    dst.sort()
    for i in range(len(dst)):
        dst[i] = arr[dst[i]]
    return dst