ステップ関数のグラフ、エラーに

Question

ステップ関数を書いてみたのですが、エラーが出てしまいます。
アドバイス、よろしくお願いいたします。

import numpy as np
import matplotlib.pyplot as plt

x=np.arange(-10,11)

def step_function(a):
    if a >0:
        return 1
    
    else:
        return 0

plt.plot(x,step_function(x))
plt.show()

◎エラー表示:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Answer 1

import numpy as np
import matplotlib.pyplot as plt

def step_function(a):
    if a >0:
        return 1
    else:
        return 0

def main():
    x=np.arange(-10,11)
    plt.plot(x,[step_function(i) for i in x])
    plt.show()
    return
    
if __name__ == '__main__':
    main()

Comments

[okahijiki]

ご回答、ありがとうございます。できましたー。for文（内包表記）にて要素を順番に代入させるのですね。