Участник:Alvant/TaskThreeEqualParts

Материал из DISCOPAL
Перейти к: навигация, поиск

https://leetcode.com/problems/three-equal-parts

import numpy as np
from typing import Callable, List, Tuple
 
 
class Solution:
    def threeEqualParts(self, A: List[int]) -> Tuple[int, int]:
        if sum(A) % 3 != 0:
            return (-1, -1)
 
        if sum(A) == 0:
            if len(A) <= 2:
                return (-1, -1)
            else:
                return (0, 2)
 
        A_as_string = ''.join(str(n) for n in A)
 
        ones_indices = np.where(np.array(A) == 1)[0]
        num_ones = sum(A)
        num_ones_in_segment = num_ones // 3
 
        i = ones_indices[num_ones_in_segment - 1]
        j = ones_indices[-num_ones_in_segment]
 
        if self.is_split_good(A_as_string, i, j):
            return (i, j)
 
        num_zeros_at_the_end = len(
            ''.join(str(n) for n in A[j:]).split('1')[-1]
        )
 
        i = self.move_pointer_to_zeros(
            i, num_zeros_at_the_end,
            lambda index: A[index] == '1' or index + 1 >= j
        )
 
        if i == -1:
            return (-1, -1)
 
        index_of_last_one_in_middle_segment = j - 1 - A[i + 1:j][::-1].index(1)
 
        k = self.move_pointer_to_zeros(
            index_of_last_one_in_middle_segment, num_zeros_at_the_end,
            lambda index: A[index] == '1' or index == j
        )
 
        if k == -1:
            return (-1, -1)
 
        j = k + 1
 
        if self.is_split_good(A_as_string, i, j):
            return (i, j)
        else:
            return (-1, -1)
 
    def is_split_good(self, A: str, i: int, j: int) -> bool:
        return (
            int(A[:i + 1], 2) == 
            int(A[i + 1:j], 2) == 
            int(A[j:], 2)
        )
 
    def move_pointer_to_zeros(self, start_index: int, num_zeros: int, is_failed_func: Callable) -> int:
        i = start_index
 
        for _ in range(num_zeros):
            i += 1
 
            if is_failed_func(i):
                return -1
 
        return i