luggage baggage

Machine learning, data analysis, web technologies and things around me.

重複あり名前リストから、スコア順に上位N人を取得しランキングを作成する(numpy.unique)

重複を含む配列があり、各々の要素に対してスコアが付随している時、スコアの大きい順に最大 N 個のユニークな要素を取得してランキングを作成する、というのを numpy.unique を使ってやる小ネタです。

例として、次のような配列を考えます。

names = ["hoge", "hoge", "var", "foo", "foo", "var"]
scores = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5]

hoge, foo, var をユニークな要素とする名前の配列があり、各々に対してスコアが付随しています。今回、前提として、両者はスコア降順でソートされているとします。

今やりたいことは、次のような配列を取得することです(N=3)。

expected_names = ["hoge", "var", "foo"]
expected_scores = [1.0, 0.8, 0.7]

これを実現するには、次のようにします。

import numpy as np


def take_top_n(names, scores, n=3):
    assert len(names) == len(scores)

    # Note: np.unique sorts the output.
    _, index = np.unique(names, return_index=True)
    sort_index = np.argsort(index)
    index = index[sort_index]

    unique_names = np.array(names)[index]
    unique_scores = np.array(scores)[index]
    top_names = unique_names[:n].tolist()
    top_scores = unique_scores[:n].tolist()

    return top_names, top_scores

主に np.unique を使うだけなのですが、この関数には少しだけ癖があるので注意が必要です。というのは、上のコードのコメントに Note と書いた通り、ユニーク要素からなる配列をソートした状態で出力してしまうのですね。今回の事例でいうと、単に np.unique しただけでは expected_names にある順序ではなく ["foo", "hoge", "var"] が返ってしまいます。そのため、np.argsort で一度配列の順序を正しくした後に上位要素を取得することが必要となります。

コードは再利用可能な形で blog/create_ranking_from_duplicated_scored_names at master · dlnp2/blog · GitHub に置きました。