Код на python не всегда работает верно

Рейтинг: 0Ответов: 2Опубликовано: 20.07.2023

я написала код по задаче:

Дан массив a длины n, вы можете сделать не более k 
операций следующего типа: выбрать 2
различных элемента в массиве, прибавить 1
к первому и вычесть 1 из второго. 
После применения операции все элемента a
должны остаться неотрицательными.
Какой лексикографически минимальный массив можно 
получить? Массив xлексикографически меньше чем 
массив y, если есть индекс i такой, что 
xi<yi, и xj=yjдля всех 1≤j<i. Проще говоря, 
для первого такого i, где массивы различны, xi<yi.

Входные данные
В первой строке записано одно целое число 
t (1≤t≤20) – количество наборов входных данных.

В первой строке каждого набора входных 
данных записаны 2целых числа n и k
(2≤n≤100, 1≤k≤10000) — размер массива 
и максимальное количество операций.

Во второй строке каждого набора входных 
данных записаны n целых чисел a1, a2, …, an
(0≤ai≤100) — элементы массива a.

Выходные данные
На каждый набор входных данных выведите 
лексикографически минимальный массив, который 
может получиться из исходного за не более чем k
операций.

Пример
входные данные
2
3 1
3 1 4
2 10
1 0
выходные данные
2 1 5 
0 1

Но на некоторых тестах он выдает неверный ответ, мой код:

n=int(input()) 
for i in range(n):
    a, b=map(int, input().split()) 
    s=list(map(int, input().split())) 
    if s.count(0)==a-1:
        print("0 "*(a-1), end="")
        print(max(s))
    else:   
        x=0
        y=a-1
        while s[x]>0 and s[y]>0 and b>0 and s.count(0)<a-1:
            s[x]-=1 
            s[y]+=1 
            b-=1 
            if s[x]==0:
                x+=1 
            if s[y]==0:
                y-=1 
        print(*s)

Не знаю что делать, помогите пожалуйста доработать и улучшить код!!!

python алгоритм
Источник: Stack Overflow на русском

Ответы

▲ 4Принят

Главная идея в вашем коде правильная: брать единички как можно левее и перемещать их как можно правее. Но есть ошибки:

  1. Тест с массивом 1 1 0 оставляет его без изменений. Хотя есть меньший ответ - 0 0 2. Причина в условии s[y]>0. Оно лишнее и только мешает. Если последний элемент массива ноль, почему мы отказываемся помещать туда единички? Убираем условие.

  2. Тест с массивом 0 1 0 1 оставляет его без изменений. Хотя меньший ответ - 0 0 0 2. Причина в условии s[x]>0. Прерывать цикл из-за нуля в текущей позиции неправильно. Справа от нуля могут быть числа, которые ещё можно обработать. Но ноль в s[x] как-то обрабатывать нужно. Поставим условие внутри цикла:

        while b>0 and s.count(0)<a-1:
            if s[x] == 0:
                x += 1 
            else:
                s[x]-=1 
                s[y]+=1 
                b-=1 
                if s[x]==0:
                    x+=1 
                if s[y]==0:
                    y-=1

  1. Условие if s[y]==0: никогда не выполняется. s[y] или уже был больше нуля или только что получил дополнительную единицу. Убираем его.

  2. Условие if s[x]==0: было полезным и осталось полезным. Но аналогичное условие уже есть в начале тела цикла. Убираем его:

        x=0
        y=a-1
        while b>0 and s.count(0)<a-1:
            if s[x] == 0:
                x += 1 
            else:
                s[x]-=1 
                s[y]+=1 
                b-=1

  1. Условие s.count(0)<a-1 полезное, но его можно заменить более простым условием - x < y. Пока x меньше y можно продолжать попытки переносить единицы, иначе цикл надо останавливать. Ещё одна причина - новое условие считается быстрее.

  2. Что делает if s.count(0)==a-1: в начале программы? Он оптимизирует ситуацию когда в массиве почти все нули. Нужно ли оно? Нет, без него можно обойтись. Программа станет медленнее (но только в специальном случае) и проще. Удаляем:

n=int(input()) 
for i in range(n):
    a, b=map(int, input().split()) 
    s=list(map(int, input().split())) 
    x=0
    y=a-1
    while b>0 and x < y:
        if s[x] == 0:
            x += 1 
        else:
            s[x]-=1 
            s[y]+=1 
            b-=1 
    print(*s)

Последняя версия стала работать правильно. Она ещё и проще. Но это не повод останавливаться.

  1. Почему мы переносим единички по одной? Можно перенести сразу min(s[x], b) единичек. Окажется что условие внутри цикла не нужно:
    while b>0 and x < y:
        m = min(s[x], b)
        s[x] -= m
        s[y] += m
        b -= m
        x += 1
  1. Условие b>0 стало не нужно. Оно делает программу немного быстрее, но усложняет. Если его удалить, цикл while можно заменить на for ... range:
n=int(input()) 
for i in range(n):
    a, b=map(int, input().split()) 
    s=list(map(int, input().split())) 
    y=a-1
    for x in range(y):
        m = min(s[x], b)
        s[x] -= m
        s[y] += m
        b -= m
    print(*s)
  1. Текст лучше отформатировать, переменные привести к терминам из задачи:
t = int(input()) 
for _ in range(t):
    n, k = map(int, input().split()) 
    x = list(map(int, input().split())) 
    for i in range(n - 1):
        m = min(x[i], k)
        x[i] -= m
        x[n - 1] += m
        k -= m
    print(*x)
▲ 2

1. Вам необходимо внести изменения в код, чтобы он соответствовал формату ввода и вывода, описанному в задаче. Для этого вам нужно изменить строку "n=int(input())" на "t=int(input())" для чтения количества наборов входных данных. Затем вы должны внести изменения в вывод результата.

2. У вас в цикле while условия проверяются в неправильном порядке. Сначала вы должны проверить, есть ли нулевые элементы в массиве, а затем уменьшать b.

3. Ваш код не удовлетворяет условию задачи, когда s.count(0) равняется a-1. В этом случае вы должны выводить весь массив равным 0.

Ниже приведен исправленный код с учетом описанных изменений:

t = int(input())
for _ in range(t):
    n, k = map(int, input().split())
    s = list(map(int, input().split()))

    if s.count(0) == n - 1:
        result = [0] * n
        result[n - 1] = max(s)
        print(*result)
    else:
        x = 0
        y = n - 1
        while b > 0 and s.count(0) < n - 1:
            if s[x] > 0 and s[y] > 0:
                s[x] -= 1
                s[y] += 1
                b -= 1
            if s[x] == 0:
                x += 1
            if s[y] == 0:
                y -= 1
        print(*s)