A function to calculate the probability of a streak of length k in n tries

streakchance.py

import functools
import sys

from argparse import ArgumentParser

@functools.cache
def streak_chance(n, k, p):
    # Probability of k successes in k tries is p^k.
    if k == n:
        return p ** k
    if k > n:
        return 0
    # New probability is probability of a streak in first n-1...
    had_streak = streak_chance(n-1, k, p)
    # Plus probability that you just got a streak with your nth try.
    # In other words, you didn't have a streak in n-k-1 tries, then had one failure, then k successes.
    new_streak = (1 - streak_chance(n-k-1, k, p)) * (1-p) * p ** k
    return had_streak + new_streak


def main():
    # Either set this to at least 2x the number of tries, or work up to the number of tries by
    # calling multiple times (e.g. streak_chance(2000, 20, 0.8), then streak_chance(4000, 20, 0.8)).
    sys.setrecursionlimit(4000);
    parser = ArgumentParser()
    parser.add_argument("tries", help="Number of runs played")
    parser.add_argument("streak", help="Desired length of streak")
    parser.add_argument("winprob", help="Probability of winning a given run")
    args = parser.parse_args()

    print(f'streak_chance({args.tries}, {args.streak}, {args.winprob}):')
    print(streak_chance(int(args.tries), int(args.streak), float(args.winprob)))


if __name__ == "__main__":
    main()