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

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_subarray (xs : List ℤ) 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

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]

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]

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
  case inv1 =>
    -- cur might be negative so the loop invariant can only guarantee ∀ i < xs.prefix.length
    exact ⇓⟨xs, ans, cur⟩ => ⌜0 ≤ ans ∧ (∃ i ≤ xs.prefix.length, (xs.prefix.drop i).sum = cur) ∧ (∀ i < xs.prefix.length, (xs.prefix.drop i).sum ≤ cur) ∧ is_max_subarray xs.prefix ans⌝
  all_goals mleave
  case vc1.step ih =>
    expose_names
    have cur_1_exists : ∃ i ≤ (pref ++ [cur]).length, (List.drop i (pref ++ [cur])).sum = cur_1 := by
      by_cases cur ≤ b.snd + cur
      · obtain ⟨i, hi, hsum⟩ := ih.2.1
        use i
        rw [drop_append_sum pref [cur] hi, hsum]
        grind
      · use pref.length
        rw [drop_append_sum pref [cur] (Nat.le_refl pref.length)]
        grind
    have cur_1_max : ∀ i < (pref ++ [cur]).length, (List.drop i (pref ++ [cur])).sum ≤ cur_1 := by
      intro i _
      rw [drop_append_sum pref [cur] (by grind)]
      by_cases i = pref.length <;> grind
    and_intros
    · grind
    · exact cur_1_exists
    · exact cur_1_max
    · by_cases b.fst ≤ cur_1
      · obtain ⟨i, _⟩ := cur_1_exists
        use i, (by grind), pref.length + 1
        constructor
        · grind
        · rw [extract_end_drop (pref ++ [cur]) (by grind)]
          grind
      · obtain ⟨i, hi, j, hj, _⟩ := ih.2.2.2.1
        use i, (by grind), j, (by grind)
        rw [extract_in_bounds pref [cur] hi hj]
        grind
    · intro i _ j _
      by_cases hi : i ≤ pref.length
      · by_cases hj : j ≤ pref.length
        · rw [extract_in_bounds pref [cur] hi hj]
          grw [ih.2.2.2.2 i hi j hj]
          exact le_max_left b.fst cur_1
        · rw [extract_end_drop (pref ++ [cur]) (by grind)]
          grw [cur_1_max i (by grind)]
          exact le_max_right b.fst cur_1
      · grind
  case vc2.a.pre => simp [is_max_subarray]
  case vc3.a.post => grind