Pythonの演算子オーバーロードに触ってみた

四則演算(+-*/)などの演算子について、クラスに独自の振る舞いを定義することができます
これは演算子オーバーロードと呼ばれます
Pythonではよくある演算子だけでなく"len()"や"del"などの様々な演算子(？)に対してもオーバーロードができます

日本語だと以下の記事が参考になりました

• とりあえず: [Python][お勉強] Python入門(50) - 演算子のオーバーロード
• お気楽 Python プログラミング入門：第５回オブジェクト指向の基礎知識

昨日の記事でも使った自作のLossy Countingのクラスには、前の記事で書いたソースコードに加えて"[]"による値の取得の演算子のオーバーロード(__getitem__)を追加してあります
こう書けばいちいち"obj.getCount(key)"とか書かなくても"obj[key]"と書けば値が得られるようになります

__getitem__と__setitem__の演算子を定義すれば"obj[key] += 1"とかも動いて驚きました(当たり前かもしれませんが

import collections

class LossyCounting(object):
    def __init__(self, epsilon):
        self.N = 0
        self.count = {}
        self.bucketID = {}
        self.epsilon = epsilon
        self.b_current = 1
    
    def getCount(self, item):
        return self.count.get(item, 0)
    
    def __getitem__(self, item):
        return self.getCount(item)

    def getBucketID(self, item):
        return self.bucketID[item]
    
    def trim(self):
        for item in self.count.keys():
            if self.count[item] <= self.b_current - self.bucketID[item]:
                del self.count[item]
                del self.bucketID[item]
        
    def add(self, item):
        self.N += 1
        if item in self.count:
            self.count[item] += 1
        else:
            self.count[item] = 1
            self.bucketID[item] = self.b_current - 1
        
        if self.N % int(1 / self.epsilon) == 0:
            self.trim()
            self.b_current += 1
    
    def iterateOverThresholdCount(self, threshold_count):
        assert threshold_count > self.epsilon * self.N, "too small threshold"
        
        self.trim()
        for item in self.count:
            if self.count[item] >= threshold_count - self.epsilon * self.N:
                yield (item, self.count[item])
    
    def iterateOverThresholdRate(self, threshold_rate):
        return self.iterateOverThresholdCount(threshold_rate * self.N)