Compare Strings

Question

You have been given two strings, A and B (of length N each) and Q queries. The strings contain only 0s and/or 1s.

For every query, you are given an index i. You have to update the value at index i to 1 in string B and check if B is lexicographically equal to or larger than A or not. If yes, then print "YES" and if not, print "NO" (without quotes).

Input format

First line contains two space-separated integers N and Q.
Next line contains the string A.
Next line contains the string B.
Next Q lines contains an integer i (1 - based indexing)

Output Format

For each query, print the desired output in a new line.

Sample Input

5 5
11111
00010
1
2
3
4
5

Sample Output

NO
NO
NO
NO
YES

Explanation

After 1st query: B = 10010. B < A. (NO)
After 2nd query: B = 11010. B < A. (NO)
After 3rd query: B = 11110. B < A. (NO)
After 4th query: B = 11110. B < A. (NO)
After 5th query: B = 11111. B = A. (YES)

My solution in Python:

# Write your code here
N_Q = list(map(int, input().split()))

string_a_list = list(input())
string_b_list = list(input())
different_set = set()
for index in range(N_Q[0]):
    if string_a_list[index] != string_b_list[index]:
        different_set.add(index + 1)

for query in range(N_Q[1]):
    index_q = int(input()) # 1 based
    index = index_q - 1
    string_b_list[index] = 1
    if len(different_set) == 0:
        break
    if int(string_a_list[index]) == int(string_b_list[index]):
        different_set.discard(index_q)
    if len(different_set) != 0:
        print("NO")


if len(different_set) == 0:
    print("YES")
if len(different_set) > 0:
    firstIndex = next(iter(different_set))
    index == firstIndex - 1
    if int(string_a_list[index]) > int(string_b_list[index]):
        print("NO")
    if int(string_a_list[index]) < int(string_b_list[index]):
        print("YES")

But for one test case that has got really very large input it's taking 5 sec so I'm getting timeout for that case:

Submission

Input #8    5.001169    22512    0 

How can I optimize this further? What are the critical points to solve problem like this that I'm missing?

Answer 1

One way to solve this is to keep a track of lexicographical ordering for every pair of suffixes. Then for every update just update the ordering for suffix starting from index i. If ordering was changed then update suffix starting from i-1 and keep on updating as long as there are changes. The current state after every query can be found from suffix[0].

n, q = (int(x) for x in input().split())
a = [int(x) for x in input()]
b = [int(x) for x in input()]

def update(idx):
    for i in range(idx, -1, -1):
        x = b[i] - a[i]
        if x == 0 and i < n - 1:
            x = suffix[i + 1]
        if x == suffix[i]:
            break
        suffix[i] = x

# Initialize to invalid value to force initial update
suffix = [-2] * n
update(n - 1)

# Buffer results for faster output
res = []
for _ in range(q):
    i = int(input()) - 1
    b[i] = 1
    update(i)
    res.append('YES' if suffix[0] >= 0 else 'NO')

print('\n'.join(res))

Since every index can be updated maximum 2 times the time complexity is O(N + Q).

Answer 2

You can
- List item
- just use list comparison, so no need to convert between int and str
- make a and b lists of 1s and 0s
- make the diff-set with a set comparison
- keep a flag. Once b is larger than or equal to a, it will stay this way

My code:

N, Q = list(map(int, input().split()))

a = list(map(int, input()))
b = list(map(int, input()))
diff = {i
    for i, (j, k) in enumerate(zip(a, b))
    if j != k
}

smaller = a > b
for _ in range(Q):
    idx = int(input()) - 1
    if smaller and idx in diff:
        b[idx] = 1
        diff.remove(idx)    
        if smaller and a > b :
            print('NO')
        else:
            print('YES')
            smaller = True
    else:
        if smaller:
            print('NO')
        else:
            print('YES')
            smaller = False

Alternative solution

What really matters is at what index the first element of A is larger than B and vice versa, so we just keep track of those, and only set B[idx] to 1 in :

N, Q = list(map(int, input().split()))

A = list(map(int, input()))
B = list(map(int, input()))

idx_a = None
idx_b = len(B)
for i, (a, b) in enumerate(zip(A, B)):
    if a > b and idx_a is None:
        idx_a = i
    if idx_a is not None and b > a:
        idx_b = i
        break
a_larger = True
if idx_a is None:
    a_larger = False

for _ in range(Q):
    idx = int(input()) - 1
    if not a_larger: 
        print('YES')
        continue
    elif idx < idx_a:
        if not B[idx]:
            print('YES')
            a_larger = False
        else:
            print('NO')
        continue
    elif idx == idx_a:
        B[idx] = 1
        idx_a = find_new_a_idx(A, B, idx_a, idx_b)
        if idx_a is None:
            print('YES')
            a_larger = False
        else:
            print('NO')
        continue
    print('NO')
    if idx >= idx_b:
        continue
        if not B[idx]:
            B[idx] = 1
            if not A[idx]:
                idx_b = idx

This is rather ugly code to take all the edge cases into account.

I tried also with array.array('b)', but that was not faster