配列の同じ要素の連続を別の要素に置き換える【Python】

概要

発話区間検出などをしていると、発話が細切れになりすぎることを避けるため、一定長さ以下の沈黙は無視したくなることがあります。そのような場合に使える、一定長さ以下/以上の同じ要素の連続を別の要素に置き換える関数を作りました。

ゴール

以下のような入出力をする関数replace_sequenceを作ることが目標です。

def replace_sequence(arr, *, before = False, after = True, min_len = 0, max_len = 5):
  """処理を書く"""

# 2以上3以下のFalseの連続をTrueに置き換える
input_arr = [False, False, True, False, False, False, False, True]
output_arr = replace_sequence(input_arr, before = False, after = True, min_len = 2, max_len = 3)
print(output_arr)
#output: [True, True, True, False, False, False, False, True]

実装

以下のように実装してみました。

def replace_sequence(arr, *, before = False, after = True, min_len = 0, max_len = 5):

  replaced_arr = [] # 置き換え後の配列
  start = -1 # 開始位置を初期化
  
  for i, v in enumerate(arr):
    # 要素がbeforeと異なり、開始位置が設定されている場合
    if v != before and start >= 0:
      # before要素の連続個数をseq_lenに代入
      seq_len = i - start 
      # seq_lenがmin_len以上max_len以下の場合、seq_len個のafterを追加
      if min_len <= seq_len <= max_len:
        replaced_arr += [after for _ in range(seq_len)]
      # それ以外の場合、seq_len個のbeforeを追加
      else:
        replaced_arr += [before for _ in range(seq_len)]
      # 要素をそのまま追加
      replaced_arr.append(v)
      # 開始位置を破棄
      start = -1
    # vがbeforeと異なり、開始位置が設定されていない場合
    elif v != before and start == -1:
      # 要素をそのまま追加
      replaced_arr.append(v)
    # vがbeforeと同じで、開始位置が設定されていない場合
    elif v == before and start == -1:
      # 開始位置を設定
      start = i
    # vがbeforeと同じで、開始位置が設定されている場合
    elif v == before and start >= 0:
      # 何もしない
      continue
  # 最後の要素がbeforeの場合（末尾にbeforeの連続がある場合）
  if start >= 0:
    seq_len = len(arr) - start
    if min_len <= seq_len <= max_len:
      replaced_arr += [after for _ in range(seq_len)]
    else:
      replaced_arr += [before for _ in range(seq_len)]
  return replaced_arr  

テストします。

import unittest

class TesteReplaceSequence(unittest.TestCase):
  def test_replace_sequence(self):
    # 2連続のFalseを置き換える
    self.assertEqual(replace_sequence(
        [True, False, False, True]
        , before = False, after = True
        , min_len = 2, max_len = 2
      ),
      [True, True, True, True]
    )
    
    # 2連続以上のFalseは置き換えない
    self.assertEqual(replace_sequence(
        [True, False, False, True]
        , before = False, after = True
        , min_len = 0, max_len = 1
      ),
      [True, False, False, True]
    )

    # TrueをFalseに置き換える
    self.assertEqual(replace_sequence(
        [True, False, False, True]
        , before = True, after = False
        , min_len = 0, max_len = 1
      ),
      [False,False, False, False]
    )


unittest.main(argv=['first-arg-is-ignored'], exit=False) # juypter notebookで実行する場合

よさそうです（テストとしては全然網羅できていないですが、おそらくは）。

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

別実装

@WolfMoon さまのコメントを参考に、rle, edit, elrの3段階による処理を実装してみました（コメントありがとうございます）。rle関数の戻り値を(value, count）のペアの配列にしていることを始め、少し、自分が理解しやすいように書き換えていますが、だいたいはそのままなので、詳細はコメントをご覧ください。

自分で思いついた実装に比べると、各段階の処理が単純なためunittestが書きやすく、保守性が優れている気がします。

class SequenceReplacer:
  @staticmethod
  def rle(arr):
    """ run length encoding
    使用法　value_counts = rle(arr)
    arr はリスト, value_counts はリストに出現する連(run)と連の長さのペアを要素とするリスト
    例
    arr = [False, False, True, False, False, False, False, True]
    rle(arr) は [(False, 2), (True, 1), (False, 4), (True, 1)] を返す
    rle('001111000000011') は [('0', 2), ('1', 4), ('0', 7), ('1', 2)] を返す
    """
  
    if len(arr) == 0:
      return []

    value_counts = []
    count = 0
    for current, next in zip(arr[:-1], arr[1:]):
      count += 1
      if current != next:
        value_counts.append((current, count))
        count = 0
    count += 1
    value_counts.append([arr[-1], count])

    return value_counts

  @staticmethod
  def edit(value_counts, before, after, min_len, max_len):
    editted_value_counts = value_counts.copy()
    for i, (value, count) in enumerate(value_counts):
      if min_len <= count <= max_len and value == before:
        editted_value_counts[i] = (after, count)
    return editted_value_counts

  @staticmethod
  def elr(value_counts):
    """ rle(arr) の逆
    使用法　arr = elr(value_counts)
    elr([(False, 2), (True, 1), (False, 4), (True, 1)]) は
    [False, False, True, False, False, False, False, True] を返す
    elr( [ ('0', 2), ('1', 4), ('0', 7), ('1', 2) ] ) は長さ 15 のリストを返すので，文字列にするには
    ''.join(elr([ ('0', 2), ('1', 4), ('0', 7), ('1', 2) ])) とすれば '001111000000011' を得る
    """
    arr = [ ]
    for value, count in value_counts:
      arr += [ value for _ in range(count) ]
    
    return arr
  
  @classmethod
  def exec(cls, arr, *, before = False, after = True, min_len = 0, max_len = 5, edit_func = None):
    """rle, edit, elrの3段階で系列の置換を実施
    入力
    arr: 置換元の配列
    before: 置換対象の要素
    after: beforeの置換後の要素
    min_len, max_len: min_len以上max_len以下のbeforeの連続を同じ長さのafterに置き換える
    edit_func: 独自ルールで編集したい場合に設定。入力・出力ともに(value, count)のリストであることが必要。
               edit_funcが設定されている場合、before, after, min_len, max_lenは無視される。設定されていない（None）場合、cls.editが使われる
    出力
    replaced_arr: 置換後の配列

    使用法
    exec([False, False, True, False, False, False, False, True], before = False, after = True, min_len = 2, max_len = 2)
    の戻り値は
    [True, True, True, False, False, False, False, True]
    （ちょうど2のFalseの連続をTrueに置換）
    """
    edit_func = edit_func or (lambda x: cls.edit(x, before, after, min_len, max_len))
    value_counts = cls.rle(arr)
    editted_value_counts = edit_func(value_counts)
    replaced_arr = cls.elr(editted_value_counts)

    return replaced_arr

# ちょうど2つのFalseの連続をTrueに置換
output = SequenceReplacer.exec([False, False, True, False, False, False, False, True], before = False, after = True, min_len = 2, max_len = 2)
print(output)

# 独自の編集ルールのテスト（結果は上に同じ）
output = SequenceReplacer.exec([False, False, True, False, False, False, False, True], edit_func = lambda x: SequenceReplacer.edit(x, False, True, 2, 2))
print(output)

[True, True, True, False, False, False, False, True]
[True, True, True, False, False, False, False, True]