Finding largest number we can make

Question

The problem is from yesterday's contest from Codeforces.

I have observed that, from the left side (i.e. MSB), for every digit, if we replace the given number with the output of f when after passing the digit - we get a bigger number we will get the desired output. I'm a beginner in programming and have written code with Python 3, which is showing TLE after submission (though it is correct and gives correct output for the sample test cases). I think a better algorithm is needed.

Variables st (start index), flag, cntr (length of contiguous subsegment ) are taken for the condition the part where we will replace the digits needs to be contiguous subsegment:

n = int(input())
s = input()
f = list(map(int,input().split()))
result = []
st = -1
flag = 1 
cntr = 0
for i in range(len(s)):
    x=f[int(s[i])-1]
    if(x>int(s[i])):
        result.append(str(x)) 
        if(flag==1):
            st = i 
            flag = 0
        if(flag==0):
            cntr+=1 
    else:
        result.append(s[i])
        if(st!=-1 and (i-st-1)!=cntr):
            break
for j in range(i+1,len(s)):
    result.append(s[j])
#print(result)
r = int(''.join(result))
print(r)

Can anyone make any improvement in my code? I know there is a tutorial on that website, but I want to learn how to develop code which I have written already.

Answer 1

The speed issue can be fixed by using PyPy 3 as it is faster (at least in most codeforces type problems, I don't know in general)

However this will result in a WA verdict. To fix this just modify the break condition as follows and it will work:

if(st!=-1 and x!=int(s[i])):

This is because the current version can sometimes break out of the loop prematurely when \$x == int(s[i])\$ because it may be better to make the segment longer.

Finally, I consider you are using too many flag variables, only one is required. Here is how I would change the code:

n = int(input())
s = input()
f = list(map(int,input().split()))
result = []
flag = 1 
for i in range(len(s)):
    x=f[int(s[i])-1]
    if(x>int(s[i])):
        result.append(str(x)) 
        if(flag==1):
            flag = 0
    else:
        result.append(s[i])
        if(flag!=1 and x!=int(s[i])):
            break
for j in range(i+1,len(s)):
    result.append(s[j])
#print(result)
r = int(''.join(result))
print(r)