LeetCode 347. Top K Frequent Elements (Python3)

はじめに

LeetCode 「347. Top K Frequent Elements」の問題を Python3 で解きました。

問題

https://leetcode.com/problems/top-k-frequent-elements/description/

問題文

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

和訳
整数配列numsと整数kが与えられた場合、最も頻繁に出現するk個の要素を返します。返される順序は任意です。

例

Input: nums = [1,1,1,2,2,3], k = 2

Output: [1,2]

Input: nums = [1], k = 1

Output: [1]

Input: nums = [1,2,1,2,1,2,3,1,3,2], k = 2

Output: [1,2]

制約

• 1 <= nums.length <= 105
• -104 <= nums[i] <= 104
• k is in the range [1, the number of unique elements in the array].
• It is guaranteed that the answer is unique.

解答1

配列の整数の最も頻繁に出現する要素を求めます。
配列をループで回しながら、要素をキーとして出現回数を辞書に登録していく方法もありますが、collections.Counterを使用することで簡潔に実装することが可能です。

Counter

collectionsモジュールに含まれるクラスで、iterableなデータの各要素の出現回数をカウントすることができます。
内部的には辞書dictが使用されており、要素がキー、出現回数がバリューとして保持しています。

most_common([n])メソッドが用意されており、引数を指定すると、最も多いn要素をカウントが多いものから順に並べたリストを返します。
今回はk個分要素を求めたいのでkを指定します。

コード

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        return [num for num, _ in collections.Counter(nums).most_common(k)]

計算量

• 時間的計算量：O(n + m log m)
  • Counter(nums)の構築：O(n)
    • 全てを走査するので
  • most_common(k):O(m log m)
    • m = ユニークな要素の数
    • 全要素のソートを行う
• 空間的計算量：O(m)
  • 要素ごとのカウントを格納

kが大きい（ユニークな要素のほとんどを取る）場合は、解答2よりも効率的です。

解答2(heapq.nlargest)

heapq.nlargest(n, iterable, key=None) を使用して実装することが可能です。
このメソッドは、iterable引数のデータセットのうち、最大値から降順にn個のリストを返します。

コード

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        counter = collections.Counter(nums)
        return [num for num, _ in heapq.nlargest(k, counter.items(), key=lambda x: x[1])]

計算量

• 時間的計算量：O(n + m log k)
  • Counter(nums)の構築：O(n)
    • 全てを走査するので
  • most_common(k):O(m log m)
    • m = ユニークな要素の数
    • k個を維持するヒープを使う
• 空間的計算量：O(m)
  • 要素ごとのカウントを格納

kが小さい場合には解答1よりも効率的です。