HackerRank Probability – Sherlock and Probability

Challenge: Sherlock and Probability

Subdomeniu: Probability (probability)

Scor cont: 20.0 / 20

Submission status: Accepted

Submission score: 1.0

Submission ID: 464723140

Limbaj: python3

Link challenge: https://www.hackerrank.com/challenges/sherlock-and-probability/problem

Cerinta

Watson gave a string $S$ to Sherlock. It is $N$ characters long and consists of only `1`s and `0`s. Now he asks:  	Given an integer $K$, I'll pick two indices $i$ and $j$ at random between $1$ and $N$, both inclusive. What's the probability that both $S[i]$ and $S[j]$ are `1` and $|i - j| \le K$?

**Input Format**   
First line contains $T$, the number of testcases. Each testcase consists of $N$(the length of $S$) and $K$ in one line and string in second line.     

**Output Format**  
Print the required probability as an irreducible fraction. If required answer is `0`, output `0/1`.   

**Constraints**   
$1 \le T \le 10^5$  
$1 \le N \le 10^5$  
$1 \le K \le N$   
$1 \le \text{Sum of N over all testcases in one file} \le 10^5$    

**Sample input**   

    2
    4 3
    1011
    4 1
    1011
    
**Sample output**   

	9/16
    5/16
    
**Explanation**    
test1: Out of 16 choices, 9 pairs of $(i,j)$ satisfy our condition.  

    (1,1), (1,3), (1,4), (3,1), (3,3), (3,4), (4,1), (4,3), (4,4)      

test2: Out of 16 choices, 5 pairs of $(i,j)$ satisfy our condition.   

    (1,1), (3,3), (4,4), (4,3), (3,4)

Cod sursa

#!/bin/python3

import math
import os
import random
import re
import sys

def solve(n, k, s):

    s= [int(x) for x in s]
    rolling_sum=sum(s[:k+1])
 
    num_pairs=0
    for i in range(k,n+k):
        if s[i-k]:
            num_pairs+=2*rolling_sum-1
            rolling_sum-=1
        if i+1<n:
            rolling_sum+=s[i+1]
    
    d=math.gcd(num_pairs, n**2)
    
    return f'{int(num_pairs/d)}/{int(n**2/d)}'
if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    t = int(input().strip())

    for t_itr in range(t):
        first_multiple_input = input().rstrip().split()

        n = int(first_multiple_input[0])

        k = int(first_multiple_input[1])

        s = input()

        result = solve(n, k, s)

        fptr.write(result + '\n')

    fptr.close()