【python】条件を満たすチョコレートのセグメントの数を求めるプログラム

自分用のメモです。

▼設問

• 正の整数が入ったlistが与えられる。(s)
• 各インデックスは板チョコのブロックを表す。
• Ronくんの誕生日をd、誕生月をmとする。
• ブロック数mで、ブロック数の合計値がdを満たす、チョコの割り方のパターン数を求める。
• ブロックは連続していること。

URL

▼sample input

s=[1,2,1,1,3,2]
d,m=3,2

▼sample output

2

▼my answer

def birthday(s, d, m): 
    ans = 0

    #指定ブロック数毎の組み合わせを作る。
    blocks = [s[i:i+m] for i in range(len(s)-0)]

    for block in blocks:
        if len(block)==m and sum(block)==d:
            ans += 1
    return ans  

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')
    n = int(input().strip())
    s = list(map(int, input().rstrip().split()))
    dm = input().rstrip().split()
    d = int(dm[0])
    m = int(dm[1])
    result = birthday(s, d, m)
    fptr.write(str(result) + '\n')
    fptr.close()

・内包表記
for文のシンプルな書き方。
式 for 変数 in イテラブル
　*イテラブル：listやset,rengeなどの連続した値を含むオブジェクト。

・スライス
イテラブル[初期値:終わり値:ステップ]

s=[1,2,1,1,3,2]
print(s[0:2])
print(s[1:4])

#
[1, 2]
[2, 1, 1]

スライスの詳細はこちら

if文の一文書きを使ってもっとシンプルにしてみる。

▼my answer(simpler)

def birthday(s, d, m): 
    blocks = [s[i:i+m] for i in range(len(s)-0)]
    return sum([(1 if len(block)==m and sum(block)==d else 0) for block in blocks])

式(True) if 条件式 else 式(False)
↓
「1 if 条件式 else 0」
真なら１を、偽なら0を返す。

※式の部分に代入は使えない
☓ ans += 1
☓ ans = ans+1