Counting swaps in a sequence sorting algorithm

Question

I have been working on this question, https://open.kattis.com/problems/sequences.

You are given a sequence, in the form of a string with characters ‘0’, ‘1’, and ‘?’ only. Suppose there are n ‘?’s. Then there are 2ⁿ ways to replace each ‘?’ by a ‘0’ or a ‘1’, giving 2ⁿ different 0-1 sequences (0-1 sequences are sequences with only zeroes and ones).

For each 0-1 sequence, define its number of inversions as the minimum number of adjacent swaps required to sort the sequence in non-decreasing order. In this problem, the sequence is sorted in non-decreasing order precisely when all the zeroes occur before all the ones. For example, the sequence 11010 has 5 inversions. We can sort it by the following moves: 11010 → 11001 → 10101 → 01101 → 01011 → 00111.

Find the sum of the number of inversions of the 2ⁿ sequences, modulo 1000000007 (10⁹+7).

Input
The first and only line of input contains the input string, consisting of characters ‘0’, ‘1’, and ‘?’ only, and the input string has between 1 to 500000 characters, inclusive.

Output
Output an integer indicating the aforementioned number of inversions modulo 1000000007

In summary, I am given a string of 1s, 0s, and ?s. All question marks are basically variables that can become 1s and 0s. In each case, I have to sum up the minimum number of inversions (a swap of two adjacent characters), needed to sort the string into non-decreasing order, i.e, all 0s before 1s. For example, 0?1? has four cases: 0010, 0011, 0110, and 0111, and need 1, 0, 2, and 0 inversions respectively, summing up to 3. I am able to get the answer right however my code is too slow:

string = input() zero = 0 # is the number of 0's total = 0 # total number of inversions questions = 0 # number of questions marks count = 0 # sum of indexes of questions marks for i in range(len(string)): if string[i] == '?': count = count + i questions = questions + 1 elif string[i] == '0': zero = zero + 1 total = total + i if questions != 0: total = total * 2 ** questions + count * 2 ** (questions - 1) triangular = 0 for i in range (zero): triangular = triangular + i Choosy = 1 for i in range(questions + 1): total = total - triangular * Choosy triangular = triangular + zero + i Choosy = Choosy * (questions - i) // (i + 1) print(total) 

Basically, the idea of my code is that for each case, the number of inversions needed is simply the sum of all indexes of 0 takeaway the (number of 0s -1)th triangular number. Therefore, all I need to do is add up all of the indexes of 0s, times by the number of cases, and add all indexes of questions marks, times the number of cases/2, since each question mark is a 0 exactly half the time. I then calculate the number of cases that within the question marks, there are 0, 1, 2 ..., (number of question mark) 0s using combinations, times by the (number of 0s -1)th triangular number, and minus this off the total. My code works kind of quickly but as I start putting around 20,000 question marks as my input, it starts to take more than a second. However, I need it to take 500,000 question marks as input and return a value within a second. Any ideas on how to make it faster?

Answer 1

Disclaimer. This ia not a proper review, but an extended comment.

Nice work with understanding the core of the problem. You just need to go one extra mile to come up with the closed form solution.

If you rewrite the loop as

Choosy = 1 sub = 0 for i in range(questions + 1): sub += triangular * Choosy triangular += zero + 1 Choosy = Choosy * (questions - i) // (i + 1) total -= sub 

it becomes obvious that the value it calculates is (\$q\$ and \$z\$ correspond to zero and questions)

\$\displaystyle \sum_{i=0}^q \binom{q}{i}T_{z+i} = \sum_{i=0}^q \binom{q}{i}\dfrac{(z+i)(z+i+1)}{2} = \dfrac{1}{2}\sum_{i=0}^q \binom{q}{i}(i^2 + (2z+1)i + z(z+1)) =\$

\$\displaystyle \dfrac{1}{2}\sum_{i=0}^q \binom{q}{i}i^2 + \dfrac{2z+1}{2}\sum_{i=0}^q \binom{q}{i}i + \dfrac{z(z+1)}{2}\sum_{i=0}^q \binom{q}{i}\$

Each term has a closed form. You don't need to loop at all.

Ironically, this problem wouldn't teach you how to iterate. Rather it teaches to not iterate.

A word of warning. You'd have to deal with quite large powers of 2. Do not compute them naively. Check out exponentiation by squaring, and take the modulo after each multiplication.

Answer 2

I don't fully understand the solution, but here are some improvements for the code:

You are computing related quantities twice:

In:

if questions != 0: total = total * 2 ** questions + count * 2 ** (questions - 1) 

you are computing 2 ** questions and 2 ** (questions - 1). You can just compute a = 2 ** (questions - 1), and then compute 2 ** questions as 2 * a. (This speed up may or may not be significant).

Bit shifting may be faster than **.

Speaking of powers of 2... you may or may not be familiar with the fact that 2 ** n == 0b1 << n. That is, raising 2 to the power n is equivalent to bit shifting binary 1 to the right n times. This is because we're dealing with binary, i.e. base-2. The same is true w.r.t. to 10 ** n being equivalent to shifting base-101 to the right n times.

Triangular numbers can be computed more efficiently using this formula:

Use the above formula to compute triangular more efficiently than:

for i in range (zero): triangular = triangular + i 

---

Style suggesitons:

• a = a + b can be replaced with a += b.
• Likewise, a = a - b can be replaced with a -= b.

for i in range(len(seq)): do_something(seq[i]) 

can be replaced with:

for a in seq: do_something(a) 

If you want to learn more about iteration in Python, this is a great resource: https://nedbatchelder.com/text/iter.html.

Answer 3

Reduce the numbers

You compute potentially huge numbers, which is slow. And you forgot about the modulo \$10^9+7\$, you're not doing that anywhere. If you do it at each step, then you work only with small numbers, which is much faster. For example, instead of 2 ** questions use pow(2, questions, 10**9 + 7).

Keeping your numbers modulo \$10^9+7\$ makes your solution \$O(n)\$ instead of \$O(n^2)\$, it's probably fast enough then.

Function

Another piece of advice is to put the computation into a function:

def inversions(string): ... return total print(inversions(input())) 

Testing against naive reference solution

That makes it more versatile, for example for testing against a naive solution in another function:

from itertools import product, combinations from operator import countOf def naive(string): inversions = 0 for bits in product(*(c if c != '?' else '01' for c in string)): inversions += countOf(combinations(bits, 2), ('1', '0')) return inversions % (10**9 + 7) def inversions(string): ... for string in '?0?', '0?1?': expect = naive(string) result = inversions(string) print(f'{string=}: {expect=}, {result=} =>', 'correct' if result == expect else 'wrong') 

(Note that my naive function uses the normal (and equivalent) definition of inversions, i.e., the number of pairs that are out of order.)

Output of that:

string='?0?': expect=3, result=3 => correct string='0?1?': expect=3, result=3 => correct 

Or more extensive testing, with all possible strings of up to nine characters (Try it online!):

passed = failed = 0 for length in range(10): for string in map(''.join, product('01?', repeat=length)): if inversions(string) == naive(string): passed += 1 else: failed += 1 print(f'Passed {passed} tests, failed {failed}.') 

Output:

Passed 29524 tests, failed 0. 

Alternative solution

I don't really understand your solution, either, here's how I solved it. I compute the answer for every prefix of the string. For that, I keep track of three values:

• strings: The number of possible strings (initially 1, doubles at every ?).
• ones: The total number of 1s in all those strings.
• inversions: The total number of inversions in all those strings.

Let's ignore the modulo for a moment, for a clearer view of the actual algorithm:

def inversions(string): strings = 1 ones = 0 inversions = 0 for c in string: if c == '0': inversions += ones elif c == '1': ones += strings else: inversions = inversions * 2 + ones ones = ones * 2 + strings strings *= 2 return inversions 

I'm using the "out-of-order pairs" method of counting again. That is, when a 0 comes at some point after a 1. So:

• Whenever I come across a 0, the number of inversions increase by how many 1s I have seen before. The number of strings stays the same, they all just get extended by a 0. So the total number of 1s also stays the same.
• Whenever I come across a 1, that doesn't cause new inversions and the number of strings also stay the same, but the number of 1s increases, one 1 for each string.
• Whenever I come across a ?, all three values are affected. The number of strings doubles, half of them get extended by a 0 and half of them by a 1. Duplicating all strings also doubles the numbers of inversions and ones in them. And then we get a few extra inversions and 1s depending on whether we extend by 0 or 1.

With modulo

Two ways to integrate the modulo (both accepted in 0.08 seconds):

With code at every calculation:

def inversions(string): mod = 10**9 + 7 strings = 1 ones = 0 inversions = 0 for c in string: if c == '0': inversions = (inversions + ones) % mod elif c == '1': ones = (ones + strings) % mod else: inversions = (inversions * 2 + ones) % mod ones = (ones * 2 + strings) % mod strings = strings * 2 % mod return inversions 

With a helper class that wraps ints, offers the operations we need, and always applies the modulo. Then the main algorithm code remains clean:

class ModInt: def __init__(self, value): self.value = value % (10**9 + 7) def __add__(self, other): return ModInt(self.value + int(other)) def __mul__(self, other): return ModInt(self.value * int(other)) def __int__(self): return self.value def __str__(self): return str(self.value) def inversions(string): strings = ModInt(1) ones = ModInt(0) inversions = ModInt(0) for c in string: if c == '0': inversions += ones elif c == '1': ones += strings else: inversions = inversions * 2 + ones ones = ones * 2 + strings strings *= 2 return inversions