turtlechanのブログ

無知の私がLinuxをいじりながら書いていくブログ

【Python】numpyのみで要素の順位付けしてみた

株の指標のRCIを計算するときに、価格の順位付け処理が必要。pandas や scipy で順位付け処理をしている例はよく見る。
今回単純な計算を numpy 以外に依存するのは嫌だと思い、numpy だけで作れないか試行錯誤して書いたスクリプトを記事にしようと思う。

※python2で書いています。

価格の順位付け

price_ranking.py
#! /usr/bin/env python
# coding: utf-8

import numpy as np


def price_ranking(closeList=[]):
    u, inv, counts = np.unique(closeList, return_inverse=True, return_counts=True)
    uniqueRankNd = np.array(np.hstack((0, counts[:-1].cumsum())), dtype='float16')
    uniqueRankNd = (counts == 1) * uniqueRankNd + (counts != 1) * (2 * uniqueRankNd + counts - 1) / counts
    priceRankNd = np.ones_like(inv) * inv.shape[0] - uniqueRankNd[inv]
    return priceRankNd


def main():
    print(price_ranking([700, 750, 780, 750, 800]))


if __name__ == '__main__':
    main()

解説ってほどではないけどざっくり。

u, inv, counts = np.unique(closeList, return_inverse=True, return_counts=True)

numpy の unique() メソッドで値の重複を取り除いてます。return_inverse に True を指定することで元の配列に戻すための配列を返してくれます。return_counts に True で配列内に重複している値の数を配列で返してくれます。便利ですね。
unique() の挙動は ユニーク要素で構成される配列:numpy.unique() が非常に分かりやすかったです。

uniqueRankNd = np.array(np.hstack((0, counts[:-1].cumsum())), dtype='float16')

同順位があった場合に同じ順位を付けてもらうための準備。
例えば [700, 750, 780, 750, 800] だった場合 [0, 1, 3, 2, 4] だと気持ち悪いので [0, 1, 3, 1, 4] になるようにしている。
hstack() メソッドはタプルで与えた配列の axis=1 方向への結合。cumsum() メソッドは累積和。
実際には現段階では [0., 1., 3., 4.]

uniqueRankNd = (counts == 1) * uniqueRankNd + (counts != 1) * (2 * uniqueRankNd + counts - 1) / counts

株のRCI指標では同順位があった場合に順位の平均で計算されるので、重複した順位の部分だけ平均順位に書き換え。

例えば、2位が2つあったなら 2位・3位 の平均 (2 + 3) / 2 = 2.5位になる。

priceRankNd = np.ones_like(inv) * inv.shape[0] - uniqueRankNd[inv]

実は今まで話していたのは、小さい方から順番に付けた順位。0点の人が一位で100点の人がビリの世界。素晴らしい。が、現実はそんなに甘くない。順位の反転をしてます。
元の配列の長さ(要素数)と同じ長さ(要素数)の配列を作って、配列の長さ(5)で初期化したものから引いているだけです。

実行結果
[5.  3.5 2.  3.5 1. ]

おわりに

あまり検証していないんで、不具合があるかもしれないです。間違い等あったらコメントで教えていただけると助かります。
似たようなことで悩んでいる方の参考になれば幸いです。

参考にしたサイト