Как реализовать функцию, возвращающую список квадратов элементов?

Что проверяют

На такой простой задаче в data/ML смотрят не на синтаксис, а на то, различаете ли вы чистый Python и векторные вычисления. Следом почти всегда спрашивают: «А если в списке десять миллионов чисел?» или «А если это колонка в pandas?».

Как отвечать

Покажите несколько уровней и скажите, когда какой уместен:

• списковое включение: [x * x for x in xs]. Идиоматично, читается сразу и работает для любого итерируемого объекта, не только list;
• map: list(map(lambda x: x * x, xs)) обычно не быстрее включения из-за вызова лямбды на каждый элемент, поэтому в Python его выбирают реже;
• генератор: (x * x for x in xs), если результат читается один раз потоком и держать весь список в памяти не нужно;
• NumPy: np.square(arr) или arr ** 2. Цикл идёт в скомпилированном коде, и на миллионах элементов это на один-два порядка быстрее чистого Python;
• pandas: s ** 2 или s.pow(2) без apply, который снова гоняет Python-функцию построчно.

Отдельно проговорите типы. Функция должна принимать и int, и float, поэтому проверка isinstance(xs, list) лишняя: она отвергает кортежи, range и генераторы. Если проверка всё же нужна, на неправильный тип бросают TypeError, а не ValueError. Добавьте аннотацию Iterable[float] -> list[float] и тест на pytest с пустым списком, отрицательными числами и нулём.

Где спотыкаются

Забывают про переполнение в NumPy. У массива с dtype=int8 или int16 квадрат молча переполнится и даст неверные числа без всякой ошибки, в отличие от питоновского int, у которого предела нет. Перед возведением стоит привести тип через astype.

Используют apply(lambda x: x ** 2) в pandas и удивляются скорости. Векторная операция на той же колонке в разы быстрее.

Путают «эффективно» с «коротко». map с лямбдой короче, но не быстрее, а экономию памяти даёт только генератор или обработка данных частями.