Need to fix search bar issue when search name in the search bar for selection

I am using antd tree component I have modified the component according to my requirement what my issue is if you see a component it has a select field when we select the field will add and also count will update but what happens Whenever will select a list after searching it from the search bar it shows the list selected count as 0 what my task is in search whatever field will appear only that field will and update or according to if user select and unselect that field then count will update like the example you can search in you can able to search in lowercase in the search bar Whenever we search a list in the search bar for some specific list name and when we select that searched sanction list then previously selected all list gets unchecked.

const { Search } = Input;
const hasSearchTerm = (n, searchTerm) =>
  n.toLowerCase().indexOf(searchTerm) !== -1;
const filterData = (arr, searchTerm) =>
  arr?.filter(
    (n) =>
      hasSearchTerm(n.title, searchTerm) ||
      filterData(n.children, searchTerm)?.length > 0
  );

function filteredTreeData(data, searchString, checkedKeys, setExpandedTree) {
  let keysToExpand = [];
  const filteredData = searchString
    ? filterData(data, searchString).map((n) => {
        keysToExpand.push(n.key);
        return {
          ...n,
          children: filterData(n.children, searchString, checkedKeys)
        };
      })
    : data;
  setExpandedTree([...keysToExpand]);
  return filteredData;
}
const MultiSelectWithTree = ({
  preExpanded = [],
  value = [],
  name = "multiSelectWithTree",
  onBlur,
  error,
  helperText,
  disabled = false,
  label,
  modalTitle,
  itemsSelectedLabel = "Items Selected",
  onChange = () => {},
  treeData,
  alterCountsCalculation = {}
}) => {
  const [expandedKeys, setExpandedKeys] = useState(preExpanded);
  const [checkedKeys, setCheckedKeys] = useState(value);
  const [tree, setTree] = useState(treeData);
  const [searchValue, setSearchValue] = useState("");
  const [checkedKeysCount, setCheckedKeysCount] = useState(0);
  const [openModal, setOpenModal] = React.useState(false);
  const handleModalClose = () => {
    setOpenModal(false);
    onChange(checkedKeys);
    setSearchValue("");
    setExpandedKeys([]);
  };
  React.useEffect(() => {
    if (searchValue) {
      const filteredData = filteredTreeData(
        treeData,
        searchValue,
        checkedKeys,
        setExpandedKeys
      );
      setTree([...filteredData]);
    } else {
      setTree(treeData);
    }
  }, [searchValue, treeData, checkedKeys]);

  React.useEffect(() => {
    setCheckedKeys(value?.length ? value : []);
  }, [value]);

  React.useEffect(() => {
    let alteredCount = 0;
    let count = checkedKeys.filter((item) => {
      if (Object.keys(alterCountsCalculation).includes(item)) {
        alteredCount = alteredCount + alterCountsCalculation[item];
        return false;
      }
      return !treeData.filter((node) => node.title == item).length;
    }).length;
    setCheckedKeysCount(count + alteredCount);
    console.log("newCount", count + alteredCount);
  }, [checkedKeys, treeData, alterCountsCalculation]);
  console.log("newCount");

  const onExpand = (expandedKeysValue) => {
    // if not set autoExpandParent to false, if children expanded, parent can not collapse.
    // or, you can remove all expanded children keys.
    setExpandedKeys(expandedKeysValue);
    //setAutoExpandParent(false);
  };

  const onCheck = (checkedKeysValue) => {
    setCheckedKeys(checkedKeysValue);
  };
  console.log("newCount");
  return (
    <>
      <FormControl
        variant="outlined"
        size="small"
        name={name}
        error={Boolean(error)}
        fullWidth
      >
        <Box alignItems="center" flexDirection="row" display="flex">
          <TextField
            label={label}
            error={Boolean(error)}
            helperText={helperText}
            onBlur={onBlur}
            name={name}
            disabled={disabled}
            fullWidth
            value={
              checkedKeysCount
                ? checkedKeysCount + " " + itemsSelectedLabel
                : ""
            }
            variant="outlined"
            InputProps={{
              readOnly: true
            }}
            size="small"
            onClick={() => (!disabled ? setOpenModal(true) : "")}
          ></TextField>
        </Box>
      </FormControl>
      <Dialog
        disableBackdropClick={true}
        onClose={handleModalClose}
        aria-labelledby="multi-select-dialog-title"
        open={openModal}
        fullWidth
        maxWidth="md"
      >
        <DialogTitle
          id="multi-select-dialog-title"
          onClose={handleModalClose}
          disableTypography
        >
          <Box
            display="flex"
            justifyContent="space-between"
            alignItems="center"
          >
            <Typography color="primary" variant="subtitle1">
              {modalTitle}
            </Typography>

            {handleModalClose ? (
              <Box
                display="flex"
                justifyContent="space-between"
                alignItems="center"
              >
                <Box mr={3}>
                  <Typography color="primary" variant="subtitle1">
                    {checkedKeysCount} {itemsSelectedLabel}
                  </Typography>
                </Box>

                <IconButton
                  aria-label="close"
                  onClick={handleModalClose}
                  size="small"
                >
                  <CloseIcon size="small" />
                </IconButton>
              </Box>
            ) : null}
          </Box>
        </DialogTitle>
        <Box position="sticky" top="0" zIndex={999} pl={5} pr={5}>
          <Search
            style={{ marginBottom: 8 }}
            placeholder="Search"
            onChange={(e) => {
              setSearchValue(e.target.value);
            }}
          />
        </Box>
        <DialogContent dividers style={{ minHeight: 150 }}>
          <Grid container>
            <Grid item xs={12}>
              <Box>
                <Tree
                  checkable
                  onExpand={onExpand}
                  titleRender={(nodeData) => {
                    const index = nodeData.title
                      .toLowerCase()
                      .indexOf(searchValue);
                    const beforeStr = nodeData.title.substr(0, index);
                    const afterStr = nodeData.title.substr(
                      index + searchValue?.length
                    );
                    const searchedString = nodeData.title.substr(
                      index,
                      searchValue?.length
                    );
                    let countSelectedChilds = 0;
                    if (nodeData.children?.length) {
                      nodeData.children.forEach((child) => {
                        if (checkedKeys?.includes(child.key)) {
                          countSelectedChilds++;
                        }
                      });
                    }
                    return index > -1 ? (
                      <span>
                        {beforeStr}
                        <span className="site-tree-search-value">
                          {searchedString}
                        </span>{" "}
                        {afterStr}{" "}
                        {countSelectedChilds ? (
                          <span>({countSelectedChilds})</span>
                        ) : (
                          ""
                        )}
                      </span>
                    ) : (
                      <span>{nodeData.title}</span>
                    );
                  }}
                  autoExpandParent={true}
                  onCheck={onCheck}
                  checkedKeys={checkedKeys}
                  expandedKeys={expandedKeys}
                  treeData={tree}
                />
              </Box>
            </Grid>
          </Grid>
        </DialogContent>
      </Dialog>
    </>
  );
};

CodeSandBoxLink