Изменять i-е следующее значение тензора каждый раз, когда значение 1 появляется в тензоре

У меня есть два тензора одинакового размера:

a = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b = [0, 1, 1, 1, 1,  1,  1,  1,  0,  1,  1,  1,  1,  1,  0,  0,  0,  1]

Тензор а имеет три области, которые отмечены последовательными значениями: область 1 — [1,2,3,4,5], область 2 — [10,11,12,13] и область 3 — [20, 21, 22, 23, 24, 25, 26, 27, 28].

Для каждого из этих регионов я хочу применить следующую логику: если одно из значений б равно 1, то следующие значения я устанавливаются равными 0. Если они уже равны 0, они остаются равными 0. После изменения значений я , ничего не происходит, пока другое значение б не станет равным 1. В этом случае следующие значения я принудительно устанавливаются в 0...

Некоторые примеры:

# i = 1

a     = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b_new = [0, 1, 0, 1, 0,  1,  0,  1,  0,  1,  0,  1,  0,  1,  0,  0,  0,  1]


# i = 2

a     = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b_new = [0, 1, 0, 0, 1,  1,  0,  0,  0,  1,  0,  0,  1,  0,  0,  0,  0,  1]


# i = 4

a     = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b_new = [0, 1, 0, 0, 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1]

Не уверен, что это поможет, но я смог разделить регионы на сегменты, выполнив следующие действия:

a_shifted = tf.roll(a - 1, shift=-1, axis=0)
a_shifted_segs = tf.math.cumsum(tf.cast(a_shifted != a, dtype=tf.int64), exclusive=True)

# a_shifted_segs = 
= [0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2]

Знаете ли вы какой-либо способ сделать это эффективно?

🤔 А знаете ли вы, что...
Python обладает обширной документацией и активным сообществом разработчиков.


4
326
3

Ответы:

Если я правильно понимаю, вы хотите в каждом разделе, определенном из списка a, сохранить первый 1, который вы встретите в b, затем обнулить следующие i элементы в этом разделе в b и снова проверить остальные элементы, если есть 1, и применить та же логика (обнуление следующих i элементов). затем перейдите к следующему разделу и т. д. Если я хорошо понимаю, как это реализовать:

def get_new_b(a, b, i):
    sect_idx = []
    start_idx = 0
    new_b = b.copy()
    for idx in range(1, len(a)):  # Find sections of array a
        if (a[idx] - a[idx-1]) != 1 or idx == len(a) - 1:
            sect_idx.append([start_idx, idx])
            start_idx = idx

    for sect_start, sect_stop in sect_idx:
        for b_idx in range(sect_start, sect_stop):
            if new_b[b_idx] == 1:
                for b_zer in range(b_idx + 1, min(b_idx + 1 + i, sect_stop)):
                    new_b[b_zer] = 0
    return new_b

за:

a = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]

b = [0, 1, 1, 1, 1,  1,  1,  1,  0,  1,  1,  1,  1,  1,  0,  0,  0, 1]

результаты будут:

print(get_new_b(a=a, b=b, i=1))
print(get_new_b(a=a, b=b, i=2))
print(get_new_b(a=a, b=b, i=4))

>>> [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1]
>>> [0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
>>> [0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1]

Решено

Вот чистый Tensorflow подход, который будет работать в режимах Eager Execution и Graph:

# copy, paste, acknowledge

import tensorflow as tf

def split_regions_and_modify(a, b, i):
  indices = tf.squeeze(tf.where(a[:-1] != a[1:] - 1), axis=-1) + 1
  row_splits = tf.cast(tf.cond(tf.not_equal(tf.shape(indices)[0], 0), 
                    lambda: tf.concat([indices, [indices[-1] + (tf.cast(tf.shape(a), dtype=tf.int64)[0] - indices[-1])]], axis=0), 
        lambda: tf.shape(a)[0][None]), dtype=tf.int32)

  def body(i, j, k, tensor, row_splits):
    k = tf.cond(tf.equal(row_splits[k], j), lambda: tf.add(k, 1), lambda: k)
    current_indices = tf.range(j + 1, tf.minimum(j + 1 + i, row_splits[k]), dtype=tf.int32)

    tensor = tf.cond(tf.logical_and(tf.equal(tensor[j], 1), tf.not_equal(j,  row_splits[k])), lambda: 
                  tf.tensor_scatter_nd_update(tensor, current_indices[..., None], tf.zeros_like(current_indices)), lambda: tensor)
    return i, tf.add(j, 1), k, tensor, row_splits 

  j0 = tf.constant(0)
  k0 = tf.constant(0)
  c = lambda i, j0, k0, b, row_splits: tf.logical_and(tf.less(j0, tf.shape(b)[0]), tf.less(k0, tf.shape(row_splits)[0]))
  _, _, _, output, _ = tf.while_loop(c, body, loop_vars=[i, j0, k0, b, row_splits])
  return output

Применение:

a = tf.constant([1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28])
b = tf.constant([0, 1, 1, 1, 1,  1,  1,  1,  0,  1,  1,  1,  1,  1,  0,  0,  0,  1])

split_regions_and_modify(a, b, 1)
# <tf.Tensor: shape=(18,), dtype=int32, numpy=array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1], dtype=int32)>

split_regions_and_modify(a, b, 2)
# <tf.Tensor: shape=(18,), dtype=int32, numpy=array([0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1], dtype=int32)>

split_regions_and_modify(a, b, 4)
# <tf.Tensor: shape=(18,), dtype=int32, numpy=array([0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1], dtype=int32)>

Здесь у вас есть решение тензорного потока, основанное на tf.scan. Я знаю, что условные операторы немного сложны, если у вас есть предложения по упрощению, я открыт для предложений. Однако, если вы знаете, как читать условные операторы, должно быть совершенно ясно, что делает код.

Здесь переменная i сообщает нам для каждой позиции в массиве, сколько еще значений b нужно перезаписать 0.

import tensorflow as tf 

a = tf.constant([1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28])
b = tf.constant([0, 1, 1, 1, 1,  1,  1,  1,  0,  1,  1,  1,  1,  1,  0,  0,  0,  1])

# Extract switches inside a
switches = tf.scan(
    lambda e, new_a: {'a': new_a, 'out': new_a != (e['a']+1)}, 
    a, 
    initializer = {'a': tf.reduce_min(a)-2, 'out': tf.constant(False)}
)['out']

# Define inputs for the scan iterations
initializer = {'b': tf.constant(False), 'i': tf.constant(0)}
elems = {'switches': switches, 'b': tf.cast(b, dtype=tf.bool)}

@tf.function
def step(last_out, new_in, max_i):
    new_i = tf.cond(
        last_out['i'] > 0, # If we are currently overwriting with 0
        lambda: tf.cond(
            new_in['switches'], # Is there a segment switch?
            lambda: tf.cond( # if switches:
                new_in['b'], # Check if b == 1
                lambda: tf.constant(max_i), # if b == 1: i = max_i
                lambda: tf.constant(0) # if b == 0: i = 0
            ),
            lambda: tf.maximum(last_out['i']-1, 0) # If no switch, decrement i
        ),
        lambda: tf.cond( # if we are currently not overwriting with 0
            new_in['b'], # check if b == 1
            lambda: tf.constant(max_i), # if b == 1: i = max_i
            lambda: tf.constant(0) # if b == 0: i = 0
        )
    )
    b = tf.cond(
        tf.equal(new_i, max_i), # Have we just reset i ?
        lambda: tf.constant(True), # If yes, we want to write b = 1
        lambda: tf.constant(False) # Otherwise, we write b = 0
    )
    
    return {'b': b, 'i': new_i}

Примеры:

outp_1 = tf.scan(lambda _last, _inp: step(_last, _inp, max_i=1), elems=elems, initializer=initializer)
print( tf.cast(outp_1['b'], tf.int32) )
# tf.Tensor([0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 0 0 1], shape=(18,), dtype=int32)

outp_2 = tf.scan(lambda _last, _inp: step(_last, _inp, max_i=2), elems=elems, initializer=initializer)
print( tf.cast(outp_2['b'], tf.int32) )
# tf.Tensor([0 1 0 0 1 1 0 0 0 1 0 0 1 0 0 0 0 1], shape=(18,), dtype=int32)

outp_4 = tf.scan(lambda _last, _inp: step(_last, _inp, max_i=4), elems=elems, initializer=initializer)
print( tf.cast(outp_4['b'], tf.int32) )
# tf.Tensor([0 1 0 0 0 1 0 0 0 1 0 0 0 0 0 0 0 1], shape=(18,), dtype=int32)

Этот ответ спонсируется lambda.