numpy에서 dict을 이용해서 값을 바꾸는 방법

특정 배열에서 dict을 이용해서 바꿔야 하는 경우가 있다.
기존에는 바꾸기 위해서 pandas에 replace를 사용하였는데, 이것이 은근 계산 비용이 많이 든다.
그래서 이 부분을 바꾸기 위해서 찾다가 발견한 것이 있어 공유한다.

Code

def replace_with_dict(ar, dic):
    # Extract out keys and values
    k = np.array(list(dic.keys()))
    v = np.array(list(dic.values()))

    # Get argsort indices
    sidx = k.argsort()

    # Drop the magic bomb with searchsorted to get the corresponding
    # places for a in keys (using sorter since a is not necessarily sorted).
    # Then trace it back to original order with indexing into sidx
    # Finally index into values for desired output.
    return v[sidx[np.searchsorted(k,ar,sorter=sidx)]]
    
    
def replace_with_dict2(ar, dic):
    # Extract out keys and values
    k = np.array(list(dic.keys()))
    v = np.array(list(dic.values()))

    # Get argsort indices
    sidx = k.argsort()

    ks = k[sidx]
    vs = v[sidx]
    return vs[np.searchsorted(ks,ar)]
    
def replace_with_dict2_generic(ar, dic, assume_all_present=True):
    # Extract out keys and values
    k = np.array(list(dic.keys()))
    v = np.array(list(dic.values()))

    # Get argsort indices
    sidx = k.argsort()

    ks = k[sidx]
    vs = v[sidx]
    idx = np.searchsorted(ks,ar)

    if assume_all_present==0:
        idx[idx==len(vs)] = 0
        mask = ks[idx] == ar
        return np.where(mask, vs[idx], ar)
    else:
        return vs[idx]

 

실험

간단하게 숫자와 문자에서 실험을 해봤다.

import pandas as pd 
import numpy as np 
dic ={334:0, 4:22, 8:31, 12:16, 16:17, 24:27, 28:18, 32:21, 36:1} 
np.random.seed(0) 
a = np.random.choice(list(dic.keys()), 20000)

 

dic = {"a": "aa","b":"bb","c":"cc","d":"dd"} 
a = np.random.choice(list(dic.keys()), 20000)

이걸 보면, 내가 평소에 자주 사용하던 Series로 변환하고 타입 변환하는 것은 엄청나게 느리다는 것을 알 수 있다.

유용한 것 같아 공유!

---

Reference

https://stackoverflow.com/questions/47171356/replace-values-in-numpy-array-based-on-dictionary-and-avoid-overlap-between-new

Replace values in NumPy array based on dictionary and avoid overlap between new values and keys