Last active
November 14, 2023 01:44
-
-
Save ryannow/7c60080b9518adad3c622043d39db861 to your computer and use it in GitHub Desktop.
A function to calculate the probability of a streak of length k in n tries
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment