1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77

import Std.Tactic.Do
import Mathlib

def kadane (A : Array ℤ) := Id.run do
  let mut cur := 0
  let mut ans := 0
  for x in A do
    cur := max x (cur + x)
    ans := max ans cur
  return ans

def is_max_nonempty_suffix (xs : List ℤ) m :=
  (∃ i ≤ xs.length, (xs.drop i).sum = m) ∧ ∀ i < xs.length, (xs.drop i).sum ≤ m

def is_max_subarray (xs : List ℤ) m :=
  0 ≤ m ∧ (∃ i ≤ xs.length, ∃ j ≤ xs.length, (xs.extract i j).sum = m) ∧ ∀ i ≤ xs.length, ∀ j ≤ xs.length, (xs.extract i j).sum ≤ m

@[grind]
lemma drop_append_sum [AddMonoid α] {xs ys : List α} (hi : i ≤ xs.length) : (xs ++ ys |>.drop i).sum = (xs.drop i).sum + ys.sum := by
  rw [List.drop_append_of_le_length hi, List.sum_append]

@[grind]
lemma extract_end_drop {xs : List α} (h : j = xs.length) : xs.extract i j = xs.drop i := by
  rw [h, List.extract_eq_drop_take, List.take_eq_self_iff, List.length_drop]

@[grind]
lemma extract_in_bounds {xs ys : List α} (hi : i ≤ xs.length) (hj : j ≤ xs.length) : (xs ++ ys).extract i j = xs.extract i j := by
  rw [List.extract_eq_drop_take, List.drop_append_of_le_length hi, List.take_append_of_le_length (by grind)]

open Std.Do

theorem kadane_correct : is_max_subarray A.toList (kadane A) := by
  generalize h : kadane A = x
  apply Id.of_wp_run_eq h
  mvcgen invariants
  · ⇓⟨xs, ans, cur⟩ => ⌜is_max_nonempty_suffix xs.prefix cur ∧ is_max_subarray xs.prefix ans⌝
  case vc1.step ih =>
    expose_names
    have hcur_1 : is_max_nonempty_suffix (pref ++ [cur]) cur_1 := by
      constructor
      · by_cases cur ≤ b.snd + cur
        · obtain ⟨i, hi, hsum⟩ := ih.1.1
          use i
          grind
        · use pref.length
          rw [drop_append_sum (Nat.le_refl pref.length)]
          grind
      · intro i _
        rw [drop_append_sum (by grind)]
        unfold is_max_nonempty_suffix at ih
        by_cases i = pref.length <;> grind
    and_intros
    · exact hcur_1.1
    · exact hcur_1.2
    · unfold is_max_subarray at ih
      grind
    · by_cases b.fst ≤ cur_1
      · obtain ⟨i, _⟩ := hcur_1.1
        use i, (by grind), pref.length + 1
        grind
      · obtain ⟨i, hi, j, hj, _⟩ := ih.2.2.1
        use i, (by grind), j, (by grind)
        grind
    · intro i _ j _
      by_cases hi : i ≤ pref.length
      · by_cases hj : j ≤ pref.length
        · rw [extract_in_bounds hi hj]
          grw [ih.2.2.2 i hi j hj]
          exact le_max_left b.fst cur_1
        · rw [extract_end_drop (by grind)]
          grw [hcur_1.2 i (by grind)]
          exact le_max_right b.fst cur_1
      · have : j - i = 0 := by grind
        rw [List.extract, this, List.take_zero]
        grind [is_max_subarray]
  case vc2.a.pre => simp [is_max_nonempty_suffix, is_max_subarray]
  case vc3.a.post => grind