python - 为每行生成特定数量的1，但仅在x =零的情况下

假设我有一个像

a = np.array([[0, 1, 1, 1, 0, 0, 0, 0, 1, 0],
             [0, 0, 1, 0, 0, 1, 0, 0, 0, 0],
             [1, 0, 0, 0, 0, 1, 0, 0, 1, 0]])

我希望每一行都有特定数量的-假设每行5个。
因此，在第一行中，我需要添加1，第二行需要3，第三行需要2。我需要在x = 0的位置随机生成那些。

我该怎么做呢？

最佳答案

这有点棘手，但这是一个完全矢量化的解决方案：

import numpy as np

def add_ones_up_to(data, n):
    # Count number of ones to add to each row
    c = np.maximum(n - np.count_nonzero(data, axis=-1), 0)
    # Make row-shuffling indices
    shuffle = np.argsort(np.random.random(data.shape), axis=-1)
    # Row-shuffled data
    data_shuffled = np.take_along_axis(data, shuffle, axis=-1)
    # Sorting indices for shuffled data (indices of zeros will be first)
    sorter = np.argsort(np.abs(data_shuffled), axis=-1)
    # Sorted row-shuffled data
    data_sort = np.take_along_axis(data_shuffled, sorter, axis=-1)
    # Mask for number of ones to add
    m = c[..., np.newaxis] > np.arange(data.shape[-1])
    # Replace values with ones or previous value depending on mask
    data_sort = np.where(m, 1, data_sort)
    # Undo sorting and shuffling
    reorderer = np.empty_like(sorter)
    np.put_along_axis(reorderer, sorter, np.arange(reorderer.shape[-1]), axis=-1)
    np.put_along_axis(reorderer, shuffle, reorderer.copy(), axis=-1)
    return np.take_along_axis(data_sort, reorderer, axis=-1)

np.random.seed(100)
data = np.array([[0, 1, 1, 1, 0, 0, 0, 0, 1, 0],
                 [0, 0, 1, 0, 0, 1, 0, 0, 0, 0],
                 [1, 0, 0, 0, 0, 1, 0, 0, 1, 0]])
n = 5
print(add_ones_up_to(data, n))
# [[0 1 1 1 1 0 0 0 1 0]
#  [0 1 1 1 0 1 0 1 0 0]
#  [1 0 0 0 0 1 1 0 1 1]]

关于python - 为每行生成特定数量的1，但仅在x =零的情况下，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/55865196/