We are given an array $arr$ of $1$s and $0$s and asked what is the probability that two randomly selected indices $i$ and $j$ satisfy
Logically, we could test every pair $(i, j)$ and check what fraction meet our criteria. Checking every pair to see if they satisfy both conditions will run in time $O(N^2)$ and will time out, so we’ll need to be more clever.
First, we check if there are any $1$s at all in our input. (This is a very special case, but we might as well check it first, since we are done if there are no $1$s.) After that check, total counts the number of pairs $(i, i)$, in other words when the two randomly chosen indices happen to be the same (and when $arr[i] = 1$).
If we have at least one $1$ in the input, what we’ll then do is set up a sliding window to track how many of the previous $k-1$ entries are $1$s as we move across the array from left to right. If the current element is a $1$, we’ll add two times the contents of the window to the running total: this is how many $1$s in the previous $k-1$ would form pairs with our current $1$ (we add two times the window, since we want to count both pairs $(i, j)$ and $(j, i)$ when $i \not= j$. We start with the window empty (window = 0) so that the first element doesn’t match with any “before” it. As we move right, we add to the window if the element is a $1$, and (once we’re at least $k$ elements in) subtract the leftmost element as it leaves the window.
Now we merely need to divide by the total possible $(i, j)$ pairs, which is just $N^2$, and reduce the fraction.
When it is possible, divide your code into functions. Here, the first one num_valid_pairs contains the logic for counting valid pairs. The solve function does the last step and produces the correct format.
We are using the fractions module to simplify the fraction in our solution.
from fractions import Fraction
def num_valid_pairs(arr, k):
# Convert the digits to integers
arr = [int(x) for x in arr]
total = sum(arr)
# If no 1s, then we're done...
if total == 0:
return 0
window = 0
for i in range(len(arr)):
if arr[i]:
total += 2*window # Either order allowed, (i,j) or (j,i)
window += 1
if i >= k and arr[i-k]:
window -= 1
return total
def solve(n, k, s):
num = num_valid_pairs(s, k)
den = len(s)**2 # len(s) squared possible pairs
if num == den:
return "1/1"
elif num:
return str(Fraction(num, den))
else:
return "0/1"