Как получить N записей из массива numpy?

Question

У меня есть следующие два numpy массива X и y:

X:

array([2, 1, 2, 1, 1, 1, 3, 2, 2, 1, 1, 3, 3, 3, 1, 1, 1, 1])

y:

array([1, 1, 1, 1, 1, 2, 3, 2, 2, 3, 3, 1, 2, 2, 4, 4, 4, 4])

Всего: 4 уникальные значения y: 1, 2, 3 и 4.

Мне нужно извлечь batch_size записей из X и y таким образом, чтобы все уникальные значения выбираются в y хотя бы один раз.

Например, если batch_size равно 6, то результат y может быть:

1, 1, 2, 3, 3, 4  # CORRECT !

Соответствующий X значения должны быть выбраны.

У меня есть следующая функция, но она может вернуть batch_s, которая, например, содержит только значения 1 и 2, но не содержит 3 и 4:

1, 1, 2, 1, 2, 1 # WRONG !

Моя функция :

def extract_batch_size(myarr, step, batch_size):
    # Function to fetch a "batch_size" amount of data from "myarr" data.

    shape = list(myarr.shape)
    shape[0] = batch_size
    batch_s = np.empty(shape)

    for i in range(batch_size):
        # Loop index
        index = ((step-1)*batch_size + i) % len(myarr)
        batch_s[i] = myarr[index]

    return batch_s

Вот как я называю эту функцию:

X_result = extract_batch_size(X, 1, 6)
y_result = extract_batch_size(y, 1, 6)

Answer 1

Вы можете сделать это так:

import numpy as np

def random_choice(arr, elem, max_num):  

    find_indices = list(filter(lambda i: arr[i]==elem, range(len(arr))))
    limit = min(max_num, len(find_indices))
    cnt = np.random.randint(1, limit + 1)
    arr_choice = np.random.choice(find_indices, cnt, replace=False)

    return arr_choice

def extract_batch_size(X, y, batch_size):

    uniq_val = np.unique(y)
    np.random.shuffle(uniq_val)
    uniq_num = len(uniq_val)

    if uniq_num == batch_size:   # too many unique values
        return X, y
    elif uniq_num < batch_size:
        indices = []
        for idx in range(uniq_num):
            cnt = batch_size - len(indices) - uniq_num + idx + 1
            find_indices = random_choice(y, uniq_val[idx], cnt)
            indices.extend(find_indices)
        return X[indices], y[indices]

    return None, None

Результат:

X = np.array([2, 1, 2, 1, 1, 1, 3, 2, 2, 1, 1, 3, 3, 3, 1, 1, 1, 1])
y = np.array([1, 1, 1, 1, 1, 2, 3, 2, 2, 3, 3, 1, 2, 2, 4, 4, 4, 4])

X_result, y_result = extract_batch_size(X, y, 6)
# indices: [17, 16, 15, 11, 6, 8]
# X_result: [1 1 1 3 3 2]
# y_result: [4 4 4 1 3 2]

Спасибо

Answer 2

С collections.Counter вы можете просто:

from collections import Counter
def extract_batch_size(myarr, batch_size):
    cnt=Counter(myarr)
    if(len(cnt.keys())>len(myarr)): #too many unique values
        return None, None
    elif(len(myarr)<batch_size): #too little array to satisfy requested batch_size
        return None, None
    elif(len(cnt.keys())==batch_size):
        res=list(cnt.keys())
        indices=[myarr.index(el) for el in res]
        return res, indices
    else:
        lst=list(cnt.keys())
        indices=[myarr.index(el) for el in lst]
        for i in cnt:
            if(cnt[i]>1):
                lst=lst+[i]*min(batch_size-len(lst),cnt[i]-1)
            if(len(lst)==batch_size):
                start=len(indices)
                for j in range(start, len(lst)):
                    k=0
                    s=myarr[k:].index(lst[j])
                    while(s in indices):
                        k=s+1
                        s=myarr[k:].index(lst[j])+k
                    indices.append(s)
                return lst, indices

Выходы для вашего случая использования:

X=[2, 1, 2, 1, 1, 1, 3, 2, 2, 1, 1, 3, 3, 3, 1, 1, 1, 1]
y=[1, 1, 1, 1, 1, 2, 3, 2, 2, 3, 3, 1, 2, 2, 4, 4, 4, 4]

X_result, X_indices= extract_batch_size(X, 6)
>>> X_result
[2, 1, 3, 2, 2, 2]

y_result, y_indices = extract_batch_size(y, 6)
>>> y_result
[1, 2, 3, 4, 1, 1]

#to output from X based on y_result indices:

>>> np.array(X)[y_indices]
[2 1 3 1 1 2]